guest-diag: sparse allocation-head scan in the pin stream

The first deltas of a 15k-page window took tens of seconds to arrive while
the client was mapping under load, so the pins still landed after the
client's cuCtxCreate deadline. The scanner now reads only the first 64 KB
of every 2 MB region - the pushes sit at the per-channel allocation heads -
which cuts the per-loop page-table work ~50x, and it logs scan timing so
the latency stays visible.

1 files changed, 43 insertions, 19 deletions

diff --git a/tools/guest-diag/guest_pin_stream.py b/tools/guest-diag/guest_pin_stream.py
index 169d048..bb3f99a 100644
--- a/tools/guest-diag/guest_pin_stream.py
+++ b/tools/guest-diag/guest_pin_stream.py
@@ -22,6 +22,8 @@ HI = int(sys.argv[4], 16) if len(sys.argv) > 4 else 0x206000000
 PAMAX = int(sys.argv[5], 16) if len(sys.argv) > 5 else 0x180000000
 INTERVAL = float(os.environ.get("STREAM_INTERVAL", "0.25"))
 BATCH = 4096
+REGION = 0x200000
+SPAN = int(os.environ.get("STREAM_SPAN", "0x10000"), 16)
 
 
 def cdt_pid():
@@ -38,28 +40,38 @@ def cdt_pid():
 
 
 def scan(pid):
-    """Present pages of the VA range as {va: pa}."""
+    """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.
+    """
     fd = os.open('/proc/%d/pagemap' % pid, os.O_RDONLY)
     out = {}
     try:
-        va = LO & ~0xFFF
-        while va < HI:
-            n = min(BATCH, (HI - va) // 0x1000)
-            if n <= 0:
-                break
-            os.lseek(fd, (va // 4096) * 8, 0)
-            data = os.read(fd, n * 8)
-            got = len(data) // 8
-            for i in range(got):
-                e = int.from_bytes(data[i * 8:i * 8 + 8], 'little')
-                if not ((e >> 63) & 1):
-                    continue
-                pfn = e & ((1 << 55) - 1)
-                pa = pfn << 12
-                if not pa or (PAMAX and pa >= PAMAX):
-                    continue
-                out[va + i * 0x1000] = pa
-            va += got * 0x1000
+        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
+                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
     finally:
         os.close(fd)
     return out
@@ -71,6 +83,7 @@ def main():
     sock = None
     seen = {}
     last_pid = None
+    last_log = 0.0
     while True:
         pid = cdt_pid()
         if pid is None:
@@ -87,13 +100,24 @@ def main():
                 except OSError:
                     sock.close()
                     sock = None
+        t0 = time.time()
         try:
             cur = scan(pid)
         except OSError:
             time.sleep(INTERVAL)
             continue
+        dt = time.time() - t0
         new = [(va, pa) for va, pa in cur.items() if seen.get(va) != pa]
         seen = cur
+        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 scan=%.3fs\n'
+                            % (time.strftime('%T'), pid, len(cur), len(new), dt))
+            except OSError:
+                pass
         if not new:
             time.sleep(INTERVAL)
             continue