#!/usr/bin/env python3
"""Minimal reproducer for the 7.1 MGLRU list_del corruption in lru_gen_del_folio().

Each worker, repeatedly: fault MEM_MB of anon memory in inside a fresh cgroup,
bounce it between two NUMA nodes with migrate_pages(2), and offline the cgroup
while that migration is still in flight.

A bottom-up search showed those three ingredients -- cross-node migration, in
flight across the offline, mapping outliving the cgroup -- are each necessary;
in particular, letting the migration finish before the rmdir does not reproduce.

Expect "list_del corruption ... but was <X>. (next=<X>)"

This is meant to produce an oops/panic, so be careful when you run it.
"""
import ctypes, mmap, os, signal, sys, time

MEM_MB = int(os.environ.get("MEM_MB", 4))
WORKERS = int(os.environ.get("WORKERS", 4))
DWELL = float(os.environ.get("DWELL", 0.3))
CG = "/sys/fs/cgroup/mglru-min"

SYS_MIGRATE_PAGES = 256
if os.uname().machine != "x86_64":
    sys.exit(f"{os.uname().machine}: syscall {SYS_MIGRATE_PAGES} is migrate_pages(2) "
             "on x86_64 only -- look up the number for this arch before running")

libc = ctypes.CDLL(None, use_errno=True)
libc.syscall.restype = ctypes.c_long
libc.syscall.argtypes = [ctypes.c_long, ctypes.c_long, ctypes.c_ulong,
                         ctypes.c_void_p, ctypes.c_void_p]


def migrate(pid, src, dst):
    """migrate_pages(2): move pid's pages from node src to node dst."""
    a, b = ctypes.c_ulong(1 << src), ctypes.c_ulong(1 << dst)
    return libc.syscall(SYS_MIGRATE_PAGES, pid, 64, ctypes.byref(a), ctypes.byref(b))


def spawn(fn, *args):
    pid = os.fork()
    if pid:
        return pid
    try:
        fn(*args)
    finally:
        os._exit(0)


def reap(pid):
    try:
        os.kill(pid, signal.SIGKILL)
        os.waitpid(pid, 0)
    except OSError:
        pass


def rmdir(path):
    try:
        os.rmdir(path)
    except OSError:
        pass


def hold(cg):
    """Join the cgroup, fault the memory in, then wait to be killed."""
    open(f"{cg}/cgroup.procs", "w").write(str(os.getpid()))
    m = mmap.mmap(-1, MEM_MB << 20, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS)
    for off in range(0, MEM_MB << 20, 4096):
        m[off] = 1
    signal.pause()


def bounce(pid, src, dst):
    """Keep migration in flight until killed."""
    while migrate(pid, src, dst) >= 0:
        src, dst = dst, src
        time.sleep(0.01)


def worker(idx, src, dst):
    want, i = int(MEM_MB * 0.9) << 20, 0
    while True:
        i += 1
        cg = f"{CG}/w{idx}-{i}"
        os.mkdir(cg)
        pid = spawn(hold, cg)
        while int(open(f"{cg}/memory.current").read()) < want:
            time.sleep(0.05)                       # until charged here
        migrate(pid, dst, src)                     # deterministic starting point
        open(f"{CG}/holding/cgroup.procs", "w").write(str(pid))
        mig = spawn(bounce, pid, src, dst)         # in flight...
        os.rmdir(cg)                               # ...across the offline
        time.sleep(DWELL)
        reap(mig)
        reap(pid)                                  # walks the stale list
        if i % 25 == 0:
            print(f"worker {idx}: {i} iterations", flush=True)


def main():
    if os.geteuid():
        sys.exit("need root")
    nodes = sorted(int(d[4:]) for d in os.listdir("/sys/devices/system/node")
                   if d.startswith("node") and d[4:].isdigit())
    if len(nodes) < 2:
        sys.exit(f"need >= 2 NUMA nodes (found {len(nodes)})")
    src = int(os.environ.get("SRC", nodes[1]))
    dst = int(os.environ.get("DST", nodes[0]))

    os.makedirs(f"{CG}/holding", exist_ok=True)
    open(f"{CG}/cgroup.subtree_control", "w").write("+memory")
    print(f"node {src} -> {dst}; {WORKERS} workers x {MEM_MB}MB", flush=True)

    kids = [spawn(worker, n, src, dst) for n in range(WORKERS)]
    try:
        os.waitpid(-1, 0)                          # they do not exit on their own
    except (KeyboardInterrupt, OSError):
        pass
    finally:
        for k in kids:
            reap(k)
        try:
            for pid in open(f"{CG}/holding/cgroup.procs").read().split():
                reap(int(pid))
        except OSError:
            pass
        time.sleep(1)
        for d in (os.listdir(CG) if os.path.isdir(CG) else []):
            rmdir(f"{CG}/{d}")
        rmdir(CG)


if __name__ == "__main__":
    main()