vgsp-sema: resolve the CE push buffer from the kernel chain (second-client stall)

The second CUDA client's GPFIFO entries in the CE push-buffer window
(0x120000000) had no injected pin: that window is never mapped into the
client's pagemap, so the scanner's window scan cannot see it, and its
only supplier - the one-shot uvbp kprobe event - is raced by two trace
clears (the scanner's pid-switch clear and semfix's trace-tail clear).
When the event loses the race the entry reads fall back to the ~250 ms
content rescue, which is the second-client stall and the same
resolution-miss class as the ISSUE-0001 lottery.

The guest scanner now walks channel->pool->manager->pushbuffer->rm_mem
in kernel memory with kcore and pins the push-buffer pages directly (no
trace event to lose), and the trace reader's dedup marks reset on a
client switch so a reused kptr cannot suppress the new client's #pb.

Boot evidence (r112): rescue reject ipin=0 for va_page=0x120023000,
rescue scans 175 ms - 1 s each, want incrementing to 39. Smoke test of
this fix on the stalled VM: 478 pin runs (~7955 pages) streamed within
seconds, first va=0x120000000 -> pa=0x108090000. A bounded kernel log
now prints the first CE push-buffer pins of each session.

1 files changed, 119 insertions, 3 deletions

diff --git a/tools/guest-diag/guest_pin_stream.py b/tools/guest-diag/guest_pin_stream.py
index 739aeca..a33a183 100644
--- a/tools/guest-diag/guest_pin_stream.py
+++ b/tools/guest-diag/guest_pin_stream.py
@@ -266,13 +266,19 @@ def _vmalloc_pages_for(kc, va):
     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)."""
+def resolve_mem_pins(kc, mem, out, default_pages=16, va_lo=0, va_hi=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.
+    """
     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)]
+    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 = kc.read64(mem + 0x60) or 0
@@ -441,6 +447,80 @@ def resolve_pb_pins(kc, pb, out):
     return resolve_mem_pins(kc, mem, out, default_pages=64)
 
 
+# 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.
+PB_CE_LO = 0x120000000
+PB_CE_HI = 0x121000000
+KC_CHAIN_BACK = 0x2000
+KC_CHAIN_SPAN = 0xA000
+KC_CHAIN_MAX = 400
+RESOLVED_KC_PB = set()
+
+
+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 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.
+    """
+    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)))
+    found = False
+    for w in cands[:KC_CHAIN_MAX]:
+        for mem in _chain_mems_at(kc, w):
+            vas = [kc.read64(mem + 0x10 + i * 8) for i in range(16)]
+            if not any(v and PB_CE_LO <= v < PB_CE_HI for v in vas):
+                continue
+            found = True
+            if mem in RESOLVED_KC_PB:
+                continue
+            RESOLVED_KC_PB.add(mem)
+            n = resolve_mem_pins(kc, mem, out, default_pages=64,
+                                 va_lo=PB_CE_LO, va_hi=PB_CE_HI)
+            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\n'
+                            % (time.strftime('%T'), mem, sem, n))
+            except OSError:
+                pass
+    return found
+
+
 # UvmGpuChannelInfo field offsets (nv_uvm_types.h, x86_64):
 #   +0x00 gpGet, +0x08 gpPut, +0x10 gpFifoEntries, +0x18 numGpFifoEntries,
 #   +0x58 gpFifoGpuVa, +0x60 gpPutGpuVa, +0x68 gpGetGpuVa.
@@ -611,6 +691,18 @@ class TraceReader:
             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.
@@ -866,6 +958,7 @@ def main():
     resolved_sem_chan = set()
     resolved_pbs = set()
     resolved_chs = set()
+    resolved_kc_pb = set()
     retry_at = {}
     tracer = TraceReader()
     kc = Kcore()
@@ -900,6 +993,7 @@ def main():
             resolved_sem_chan = set()
             resolved_pbs = set()
             resolved_chs = set()
+            resolved_kc_pb = set()
             retry_at = {}
             # Client switch: the old client's channel/put pointers are freed
             # memory. Clearing the #put map stops stale reads from being
@@ -907,6 +1001,8 @@ def main():
             CHAN_PUTS.clear()
             PUT_LAST.clear()
             clear_trace()
+            tracer.reset()
+            RESOLVED_KC_PB.clear()
             if sock is not None:
                 try:
                     sock.sendall(("#pid %d\n" % pid).encode())
@@ -987,6 +1083,26 @@ def main():
                     added += len(got)
                 else:
                     retry_at[("cy", sem)] = now_r + 2.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 above.
+                if sem in resolved_kc_pb or \
+                        now_r < retry_at.get(("pk", sem), 0):
+                    continue
+                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
             for pb in list(seen_pbs):
                 if pb in resolved_pbs or now_r < retry_at.get(pb, 0):
                     continue