#!/usr/bin/env python3
# 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 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 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")
SEM_EVERY = int(os.environ.get("STREAM_SEM_EVERY", "4"))


class TraceReader:
    """Tail the guest trace for new uvmt sems / uvbp pushbuffers in a thread.

    Reading the trace can block for seconds under client load; doing it in
    the scan loop delayed the first GR-window delta (observed a 3 s gap
    right when the client's channels appeared). The thread appends to a
    queue the main loop drains without blocking.

    Push-buffer kptrs are emitted as soon as the client begins its first
    push: the host resolves them through the UVM chain and pins the pages,
    which is what makes the first CE channel's releases applicable without
    waiting for the content sweep.
    """

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

    def _run(self):
        while not self.stop:
            try:
                out = subprocess.run(["tail", "-c", "262144", TRACE],
                                     capture_output=True, timeout=10).stdout
            except Exception:
                time.sleep(1)
                continue
            with self.lock:
                for m in re.finditer(rb"sem=0x([0-9a-fA-F]+)", out):
                    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]+)", out):
                    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)
            time.sleep(0.3)

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


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

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


def cdt_pid():
    prefixes = tuple(p for p in os.environ.get("CLIENT_PREFIXES",
                                               "cdt,va").split(",") if p)
    for d in os.listdir('/proc'):
        if not d.isdigit():
            continue
        try:
            with open('/proc/%s/comm' % d) as f:
                name = f.read().strip()
                if any(name.startswith(p) for p in prefixes):
                    return int(d)
        except OSError:
            pass
    return None


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
            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_vmas(fd, pid, out):
    """Present pages of the client's high-VA nvidia//dev/zero mappings.

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


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

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

    Two windows: the GR window (rings/pushbuffers/counter page) and the CE
    push-buffer window (0x120000000+, one 1 MB allocation per CE channel).
    The CE window uses a wider span because the per-channel allocations are
    1 MB apart inside a 2 MB region.
    """
    fd = os.open('/proc/%d/pagemap' % pid, os.O_RDONLY)
    out = {}
    scan.loop += 1
    try:
        scan_window(fd, LO, HI, SPAN, out)
        if HI2 > LO2:
            scan_window(fd, LO2, HI2, SPAN2, out)
        if 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()
    tracer = TraceReader()
    last_pid = None
    last_log = 0.0
    loop = 0
    while True:
        pid = cdt_pid()
        loop += 1
        if pid is None:
            last_pid = None
            seen = {}
            seen_sems = set()
            time.sleep(INTERVAL)
            continue
        if pid != last_pid:
            last_pid = pid
            seen = {}
            seen_sems = set()
            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]
        seen = cur
        lines = []
        for ctl in tracer.drain():
            if ctl.startswith("#sem"):
                seen_sems.add(int(ctl.split()[1], 16))
            lines.append(ctl)
        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:
            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)
        except OSError:
            sock.close()
            sock = None
        time.sleep(INTERVAL)


if __name__ == '__main__':
    main()