#!/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 socket
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)


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 = {}
    last_pid = None
    last_log = 0.0
    while True:
        pid = cdt_pid()
        if pid is None:
            last_pid = None
            seen = {}
            time.sleep(INTERVAL)
            continue
        if pid != last_pid:
            last_pid = pid
            seen = {}
            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
        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
        new.sort()
        lines = []
        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()