Implement proof of concept

3 files changed, 178 insertions, 0 deletions

diff --git a/pyproject.toml b/pyproject.toml
index ce85a55..dd56c1d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,6 +18,7 @@ dependencies = [
 	"brotli",
 	"pycapnp",
 	"setuptools",
+	"python-ptrace",
 	"cpuid @ git+https://github.com/taugoust/cpuid.py.git@master",
 ]
 
diff --git a/run.py b/run.py
new file mode 100644
index 0000000..870b23a
--- /dev/null
+++ b/run.py
@@ -0,0 +1,166 @@
+import os
+import signal
+import socket
+import subprocess
+import sys
+
+import ptrace.debugger
+from ptrace.debugger import (
+    ProcessSignal,
+    ProcessExit,
+    NewProcessEvent,
+)
+
+
+# -------- scheduler via unix socket --------
+
+def read_next_tid(sock, processes):
+    """
+    Read the next thread ID to schedule from a controller over a socket.
+    Blocks until a valid TID is received.
+    """
+    while True:
+        data = sock.recv(64)
+        if not data:
+            continue  # ignore empty reads
+
+        try:
+            tid = int(data.strip())
+        except ValueError:
+            print(f"Invalid scheduler value: {data!r}")
+            continue
+
+        if tid in processes:
+            print(f"Scheduler selected TID {tid}")
+            return processes[tid]
+
+        print(f"TID {tid} not active, waiting for a valid one …")
+
+
+# -------- wait until first clone --------
+
+def run_until_first_clone(debugger, process):
+    process.cont()
+
+    while True:
+        try:
+            sig = debugger.waitSignals()
+            sig.process.cont(sig.signum)
+
+        except NewProcessEvent as event:
+            new_proc = event.process
+            print(f"First clone: TID {new_proc.pid}")
+            return
+
+        except ProcessExit as event:
+            print(f"Process {event.process.pid} exited before cloning")
+            raise
+
+
+# -------- main tracer --------
+
+def trace(pid, sched_socket_path):
+    debugger = ptrace.debugger.PtraceDebugger()
+
+    debugger.traceClone()
+    debugger.traceFork()
+    debugger.traceExec()
+
+    print(f"Attach process {pid}")
+    proc0 = debugger.addProcess(pid, False)
+
+    # Create listening scheduling socket
+    if os.path.exists(sched_socket_path):
+        os.unlink(sched_socket_path)
+
+    srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    srv.bind(sched_socket_path)
+    srv.listen(1)
+
+    print(f"Waiting for scheduler connection on {sched_socket_path}")
+    conn, _ = srv.accept()
+    print("Scheduler connected")
+
+    # 1) run until clone
+    run_until_first_clone(debugger, proc0)
+
+    # 2) attach all existing threads after clone
+    for tid_str in os.listdir(f"/proc/{pid}/task"):
+        tid = int(tid_str)
+        if tid not in debugger.dict:
+            try:
+                debugger.addProcess(tid, False)
+            except Exception as e:
+                print(f"Failed to attach TID {tid}: {e}")
+
+    # 3) main loop: scheduler picks TID → run it to next syscall
+    while len(debugger.list) != 0:
+        proc = read_next_tid(conn, debugger.dict)
+        tid = proc.pid
+
+        try:
+            proc.syscall()         # continue until syscall entry/exit
+            proc.waitSyscall()
+
+            ip = proc.getInstrPointer()
+            print(f"TID {tid} syscall-stop at {hex(ip)}")
+
+        except ProcessSignal as ev:
+            ev.process.cont(ev.signum)
+
+        except ProcessExit as ev:
+            print(f"Thread {ev.process.pid} exited (exitcode={ev.exitcode})")
+            try:
+                ev.process.detach()
+            except Exception:
+                pass
+            debugger.deleteProcess(ev.process)
+
+        except Exception as e:
+            print(f"TID {tid} exception: {e}")
+            try:
+                proc.detach()
+            except Exception:
+                pass
+            debugger.deleteProcess(proc)
+
+        # detect new threads
+        for p in list(debugger.list):
+            t = p.pid
+            if t not in debugger.dict:
+                debugger.dict[t] = p
+
+    conn.close()
+    srv.close()
+    debugger.quit()
+
+
+# -------- entry point --------
+
+def quoted(s: str) -> str:
+    return f'"{s}"'
+
+
+if __name__ == "__main__":
+    env = os.environ.copy()
+
+    qemu = [
+        "qemu-x86_64",
+        "/nix/store/dmpq06y392i752zwhcna07kb2x5l58l5-memcached-static-x86_64-unknown-linux-musl-1.6.37/bin/memcached",
+        "-p", "11211",
+        "-t", "4",
+        "-vv",
+    ]
+
+    sched_path = "/tmp/memcached_scheduler.sock"
+
+    proc = subprocess.Popen(qemu, env=env)
+    try:
+        trace(proc.pid, sched_path)
+    except Exception as e:
+        print(f"Got exception: {e}")
+        proc.kill()
+        exit(2)
+
+    exit(0)
+
diff --git a/uv.lock b/uv.lock
index dc00fbe..8bb04ed 100644
--- a/uv.lock
+++ b/uv.lock
@@ -90,6 +90,7 @@ dependencies = [
     { name = "miasm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
     { name = "orjson", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
     { name = "pycapnp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "python-ptrace", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
     { name = "setuptools", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
 ]
 
@@ -112,6 +113,7 @@ requires-dist = [
     { name = "pycapnp" },
     { name = "pyright", marker = "extra == 'dev'" },
     { name = "pytest", marker = "extra == 'dev'" },
+    { name = "python-ptrace" },
     { name = "ruff", marker = "extra == 'dev'" },
     { name = "setuptools" },
 ]
@@ -297,6 +299,15 @@ wheels = [
 ]
 
 [[package]]
+name = "python-ptrace"
+version = "0.9.9"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/a5/8f/9700154c92854687c78380a5fa226233d89b2631c74e420b25c93ebc6d95/python-ptrace-0.9.9.tar.gz", hash = "sha256:56bbfef44eaf3a77be48138cca5767cdf471e8278fe1499f9b72f151907f25cf", size = 108964, upload-time = "2024-03-13T14:57:23.332Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/3e/33/0865143ddf633d27881296e77bd66b24df5faa9d744267c1d29aeb902a07/python_ptrace-0.9.9-py2.py3-none-any.whl", hash = "sha256:8c97d23d55e551e7d7c5b104b87635fdd27d6c188c9cc205e8d0f6d71272b712", size = 104763, upload-time = "2024-03-13T15:00:01.045Z" },
+]
+
+[[package]]
 name = "pytokens"
 version = "0.2.0"
 source = { registry = "https://pypi.org/simple" }