#!/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 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.25"))
BATCH = 4096
REGION = 0x200000
SPAN = int(os.environ.get("STREAM_SPAN", "0x10000"), 16)
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"))


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():
    for d in os.listdir('/proc'):
        if not d.isdigit():
            continue
        try:
            with open('/proc/%s/comm' % d) as f:
                if f.read().strip() == 'cdt':
                    return int(d)
        except OSError:
            pass
    return None


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.
    """
    fd = os.open('/proc/%d/pagemap' % pid, os.O_RDONLY)
    out = {}
    try:
        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


def main():
    with open('/tmp/guest_pin_stream.pid', 'w') as f:
        f.write('%d\n' % os.getpid())
    sock = None
    seen = {}
    seen_sems = set()
    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
        new = [(va, pa) for va, pa in cur.items() if seen.get(va) != pa]
        seen = cur
        lines = []
        if loop % SEM_EVERY == 0:
            for sem in new_sems(seen_sems):
                lines.append('#sem 0x%x\n' % sem)
        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\n'
                            % (time.strftime('%T'), pid, len(cur), len(new),
                               len(seen_sems), dt))
            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()