From 55e9d9fdca4baafce728659d43614e203bfbea9a Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 15:02:05 +0200 Subject: codemode: close the performance gap with the old worker (heads 35x -> 2.2x) Two changes, both about work that was never ours to do, found by profiling the A/B benchmark rather than guessing. 1. Serialise inside the database process. Code Mode runs to_jsonable() over whatever a snippet returns, walking the entire structure to make it JSON-safe. Our answers are already JSON-safe and they are large: a 200-row listing page is ~10k small objects, and walking them cost 66ms of the page's 92ms -- 114x what json.dumps of the very same data costs (0.58ms). Snippets now return one pre-serialised string, so that walk is O(1) and the client parses a payload it was going to parse anyway. heads(200): 92ms -> 24.7ms. 2. Detach the runtime's trace hook while our snippet runs. ida_codemode.runtime wraps every execute_python in sys.settrace(timeout_trace) to enforce deadlines, and timeout_trace RETURNS ITSELF -- which switches on LINE tracing in every frame it sees. Every line of every function we call pays a Python-level callback. Measured here: ida_bytes.get_flags 0.106us untraced 5.49us traced 52x (plain idalib, no Code Mode: 0.119us -- i.e. untraced == native) heads(200 rows) 2.0ms untraced 20.2ms traced 10x That one hook was the entire residual gap against the old unix-socket worker. The snippet now detaches it and restores it in a finally. What that gives up, stated plainly: the deadline is no longer enforced for a pure-Python loop inside our snippet. The runtime's other cancellation path -- a threading.Timer calling ida_kernwin.set_cancelled() -- does not go through the trace and still fires, so a long IDA operation remains interruptible, and every operation here is bounded by its own count/limit argument. Set IDATUI_CODEMODE_TRACE=1 to keep the stock behaviour. Against the worker backend, same box, targets/echo (worker -> codemode): heads_200 2.65ms -> 5.85ms 2.2x (was 35x) heads_500 6.21ms -> 10.67ms 1.7x heads expect-hit 2.16ms -> 4.61ms 2.1x disasm_200 9.03ms -> 5.25ms 0.6x faster decompile_cold 162.62ms -> 30.66ms 0.2x faster decompile_warm 30.02ms -> 25.39ms 0.8x faster decomp_map 45.82ms -> 47.85ms 1.0x parity pc_nums 19.90ms -> 22.56ms 1.1x parity rename_func 254.08ms -> 255.54ms 1.0x parity connect 550.0ms -> 410.0ms 0.7x faster What is left is the transport floor: an empty execute_python round trip is 2.0ms, so trivial calls (data_type 0.07ms -> 2.63ms, force_recompile, a single xref query) look like 40x while being 2.5ms of wall clock. Reducing those needs fewer calls, not faster ones -- the digest/expect path already does that for the listing, which is where call volume actually is. Full suite: 788 passed, 0 failed, 115.3s (was 146.3s; the pilot alone went 80.9s -> 62.2s). --- idatui/codemode_client.py | 71 +++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 66 insertions(+), 5 deletions(-) diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 9d91335..88d7b31 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -22,7 +22,7 @@ import shlex import threading import time from pathlib import Path -from textwrap import dedent +from textwrap import dedent, indent from typing import Any from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session @@ -165,10 +165,63 @@ def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]: return processor, loading_address, file_type +#: Key of the pre-serialised payload envelope. See _script(). +_PACKED = "__idatui_json__" + +#: Serialise the answer INSIDE the database process and hand back one string. +#: +#: Code Mode runs to_jsonable() over whatever a snippet returns, walking the +#: whole structure to make it JSON-safe. Our answers are already JSON-safe, and +#: they are big: a 200-row listing page is ~10k small objects, which costs 66ms +#: to walk -- 72% of the page's total cost, and 114x what json.dumps of the very +#: same data costs (0.58ms). Returning a STRING makes that walk O(1); the client +#: parses it, which it was going to do at the transport layer anyway. +_PACK_EPILOGUE = ( + '\n{"' + _PACKED + '": json.dumps(result, separators=(",", ":"), default=str)}\n' +) + + +#: Keep Code Mode's per-line trace hook installed while our snippet runs. +#: Set IDATUI_CODEMODE_TRACE=1 to restore the stock behaviour. +_KEEP_TRACE = os.environ.get("IDATUI_CODEMODE_TRACE", "") not in ("", "0") + + def _script(args: dict[str, Any], body: str) -> str: - """Bind JSON arguments without interpolating user text into Python code.""" + """Bind JSON arguments without interpolating user text into Python code. + + Also runs the body with Code Mode's trace hook detached, which is worth an + order of magnitude. The runtime wraps every execute_python in + sys.settrace(timeout_trace), and that trace function RETURNS ITSELF, which + turns on line tracing in every frame it sees -- so every line of every + function we call pays a Python-level callback. Measured on this box: + ida_bytes.get_flags is 0.106us untraced (0.119us in a plain idalib process) + and 5.49us traced, 52x; a 200-row listing page is 2.0ms untraced and 20.2ms + traced. That single hook was the whole residual gap against the old worker. + + What this gives up: the deadline is no longer enforced for a pure-Python + loop inside our snippet. The runtime's OTHER cancellation path -- a + threading.Timer that calls ida_kernwin.set_cancelled() -- is independent of + the trace and still fires, so a long IDA operation is still interruptible; + and every operation here is bounded by its own count/limit argument. The + trace is restored in a finally, so a raising snippet cannot leak the change. + """ encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":")) - return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n" + head = f"import json\na = json.loads({encoded!r})\n" + if _KEEP_TRACE: + return f"{head}{dedent(body).strip()}\n{_PACK_EPILOGUE}" + return ( + f"{head}" + "import sys\n" + "_idatui_trace = sys.gettrace()\n" + "sys.settrace(None)\n" + "try:\n" + f"{indent(dedent(body).strip(), ' ')}\n" + ' _idatui_packed = {"' + _PACKED + '": json.dumps(' + 'result, separators=(",", ":"), default=str)}\n' + "finally:\n" + " sys.settrace(_idatui_trace)\n" + "_idatui_packed\n" + ) _OPERATIONS: dict[str, str] = { @@ -1179,6 +1232,13 @@ class CodeModeClient: raise IDAToolError("execute_python", "Code Mode returned an invalid execution result") return response["result"] + @staticmethod + def _unpack(answer: Any) -> Any: + """Undo _PACK_EPILOGUE. Anything else passes through untouched.""" + if isinstance(answer, dict) and _PACKED in answer: + return json.loads(answer[_PACKED]) + return answer + def invoke(self, operation: str, *, timeout: float | None = None, **args) -> Any: """Execute one TUI domain operation through Code Mode.""" if operation in ("idb_save", "save"): @@ -1189,12 +1249,13 @@ class CodeModeClient: if body is None: raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}") try: - answer = self.execute_python(_script(args, body), timeout=timeout) + answer = self._unpack(self.execute_python(_script(args, body), timeout=timeout)) if isinstance(answer, dict) and answer.get(_NEED_LIB): # First call against this database process (or a restarted one). self.execute_python(_script({"source": _REMOTE_LIB}, _INSTALL_LIB), timeout=timeout) - answer = self.execute_python(_script(args, body), timeout=timeout) + answer = self._unpack( + self.execute_python(_script(args, body), timeout=timeout)) return answer except IDAToolError as exc: if exc.tool == "execute_python": -- cgit v1.3.1-sl0p