diff options
| -rw-r--r-- | pyproject.toml | 1 | ||||
| -rw-r--r-- | run.py | 166 | ||||
| -rw-r--r-- | uv.lock | 11 |
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" } |