guest-diag: content-score CE mem candidates and refresh stale pin sets

The kptr-keyed one-shot latches assumed object identity, but UVM frees and
re-creates CE push-buffer rm_mems and channels mid-session and the slab
reuses the same addresses; a resolution taken before such a move pinned
pages that read all-zero when the host used them, so every CE ring entry
stalled in the 8 s grace (r125: the live pushes sat at new vmap pages of
the same rm_mem kptr). Candidates are now scored by their window content
(a live push buffer is full of method streams, a freed range reads zeros)
and the best is pinned - all-zero candidate sets emit nothing instead of
poisoning the host map. A resolved mem is re-validated every ~2 s against
its own {cpu,gpu} fields and first-page PA; when they move, the latches
drop and the current truth is re-emitted.

1 files changed, 174 insertions, 42 deletions

diff --git a/tools/guest-diag/guest_pin_stream.py b/tools/guest-diag/guest_pin_stream.py
index 1da78fb..b031733 100644
--- a/tools/guest-diag/guest_pin_stream.py
+++ b/tools/guest-diag/guest_pin_stream.py
@@ -541,12 +541,26 @@ def resolve_pb_pins(kc, pb, out):
 # 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 fingerprint (freshness)
+KC_REFRESH_S = 2.0		# re-validate a resolved CE mem at most this often
 
 
 def _chain_mems_at(kc, w):
@@ -567,13 +581,86 @@ def _chain_mems_at(kc, w):
     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
+    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.
+    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
@@ -582,29 +669,50 @@ def resolve_pb_chain(kc, sem, out):
         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
+    mems = []
     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)
-            off = mem_page_off(kc, mem)
-            n = resolve_mem_pins(kc, mem, out, default_pages=64,
-                                 va_lo=PB_CE_LO, va_hi=PB_CE_HI,
-                                 page_off=off)
-            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 size=%d\n'
-                            % (time.strftime('%T'), mem, sem, n, off,
-                               _mem_size(kc, mem)))
-            except OSError:
-                pass
-    return found
+            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)
+    off = mem_page_off(kc, mem)
+    n = resolve_mem_pins(kc, mem, out, default_pages=64,
+                         va_lo=PB_CE_LO, va_hi=PB_CE_HI,
+                         page_off=off)
+    st = {"mem": mem, "gpu": gpu, "cpu": cpu,
+          "page_va": cpu, "page_pa": kc.va_to_pa(cpu) or 0, "page_off": 0}
+    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 size=%d score=%d\n'
+                    % (time.strftime('%T'), mem, sem, n, off,
+                       _mem_size(kc, mem), _score))
+    except OSError:
+        pass
+    return True
 
 
 # UvmGpuChannelInfo field offsets (nv_uvm_types.h, x86_64):
@@ -1040,6 +1148,7 @@ def main():
     resolved_chs = set()
     resolved_kc_pb = set()
     retry_at = {}
+    kc_refresh_at = {}
     tracer = TraceReader()
     kc = Kcore()
     kc_ok = kc.load()
@@ -1075,6 +1184,7 @@ def main():
             resolved_chs = set()
             resolved_kc_pb = set()
             retry_at = {}
+            kc_refresh_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.
@@ -1083,6 +1193,7 @@ def main():
             clear_trace()
             tracer.reset()
             RESOLVED_KC_PB.clear()
+            KC_PB_STATE.clear()
             if sock is not None:
                 try:
                     sock.sendall(("#pid %d\n" % pid).encode())
@@ -1144,6 +1255,47 @@ def main():
                         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:
+                        _st = KC_PB_STATE.pop(sem, None)
+                        if _st and _st.get("mem"):
+                            RESOLVED_KC_PB.discard(_st["mem"])
+                        resolved_kc_pb.discard(sem)
+                        retry_at.pop(("pk", sem), None)
+                        dbg("kc-chain stale: re-resolving sem=0x%x" % sem)
+                # 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
@@ -1163,26 +1315,6 @@ 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