diff options
| -rw-r--r-- | idatui/pool.py | 204 | ||||
| -rw-r--r-- | idatui/worker_client.py | 27 | ||||
| -rw-r--r-- | tests/test_pool.py | 163 |
3 files changed, 388 insertions, 6 deletions
diff --git a/idatui/pool.py b/idatui/pool.py new file mode 100644 index 0000000..df852f8 --- /dev/null +++ b/idatui/pool.py @@ -0,0 +1,204 @@ +"""WorkerPool — keeps a live idalib worker per project binary, within a budget. + +One worker process holds exactly one database (idalib is single-DB and +main-thread-only), so a project with N binaries means up to N processes. They are +not cheap and they do not share: a worker on ``bash`` measures ~126 MB RSS / +117 MB PSS, and the database working set dominates for anything larger +(``libcrypto.so.3``'s ``.i64`` alone is 72 MB). + +Residency is therefore bounded by a **memory budget**, not a worker count — a +count is the wrong knob when one project holds both a 50 KB helper and a 6 MB +crypto library. Workers are spawned lazily on first use, kept resident while they +fit, and least-recently-used ones evicted when they don't. Eviction **saves the +database first**, so coming back is a load rather than a re-analysis. + +The pool never evicts the active binary, nor anything pinned. +""" +from __future__ import annotations + +import os + +from .project import BinaryRef, Project + +#: Fallback budget if /proc/meminfo can't be read (MB). +_FALLBACK_BUDGET_MB = 2048 + + +def _total_ram_mb() -> int: + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemTotal:"): + return int(line.split()[1]) // 1024 + except (OSError, ValueError, IndexError): + pass + return 0 + + +def _pss_mb(pid: int | None) -> int: + """Proportional set size of a worker, in MB. + + PSS (not RSS) is the honest per-worker cost: it splits shared pages between + the processes mapping them. In practice workers share very little, so the two + are close, but PSS is what makes summing across workers meaningful. + """ + if not pid: + return 0 + try: + with open(f"/proc/{pid}/smaps_rollup") as f: + for line in f: + if line.startswith("Pss:"): + return int(line.split()[1]) // 1024 + except (OSError, ValueError, IndexError): + pass + return 0 + + +def _default_spawn(ref: BinaryRef, ttl: int): # pragma: no cover - needs idalib + from .worker_client import WorkerClient + return WorkerClient(ref.staged, ttl=ttl) + + +class WorkerPool: + """Live workers for a project's binaries, keyed by label.""" + + def __init__(self, project: Project, *, budget_mb: int | None = None, + ttl: int = 1800, spawn=None, mem_fn=None) -> None: + self.project = project + self._ttl = ttl + self._spawn = spawn or _default_spawn + self._mem = mem_fn or (lambda c: _pss_mb(getattr(c, "pid", None))) + self._clients: dict[str, object] = {} + self._lru: list[str] = [] # least-recently-used first + self._pinned: set[str] = set() + self.active: str | None = None # never evicted + if budget_mb is None: + ram = _total_ram_mb() + budget_mb = (ram * project.memory_pct // 100) if ram else _FALLBACK_BUDGET_MB + self.budget_mb = max(budget_mb, 256) + self.evicted: list[str] = [] # labels evicted, most recent last + + # -- residency --------------------------------------------------------- # + def resident(self) -> list[str]: + """Labels with a live worker, least-recently-used first.""" + return list(self._lru) + + def is_resident(self, label: str) -> bool: + return label in self._clients + + def memory_mb(self) -> int: + return sum(self._mem(c) for c in self._clients.values()) + + def pin(self, label: str, on: bool = True) -> None: + """Keep ``label`` resident regardless of the budget.""" + self._pinned.add(label) if on else self._pinned.discard(label) + + def is_pinned(self, label: str) -> bool: + return label in self._pinned + + # -- acquire ----------------------------------------------------------- # + def get(self, label: str, progress=None): + """A live client for ``label``, spawning it (and making room) if needed. + + Staging and the scratch sweep happen here: a worker killed hard last time + leaves unpacked ``.id0/.id1/...`` behind, and the database then refuses to + reopen. Nothing else holds this DB (one worker per label), so it is safe. + """ + client = self._clients.get(label) + if client is not None: + self._touch(label) + return client + ref = self.project.by_label(label) + if ref is None: + raise KeyError(f"no such binary in the project: {label}") + + def note(msg: str) -> None: + if progress is not None: + progress(msg) + + note(f"staging {ref.label}\u2026") + self.project.stage(ref) + self.project.sweep_scratch(ref) + note(f"opening {ref.label}\u2026") + client = self._spawn(ref, self._ttl) + connect = getattr(client, "connect", None) + if connect is not None: + connect(progress=progress) if progress is not None else connect() + self._clients[label] = client + self._lru.append(label) + self._enforce_budget(protect=label) + return client + + def _touch(self, label: str) -> None: + if label in self._lru: + self._lru.remove(label) + self._lru.append(label) + + def set_active(self, label: str | None) -> None: + self.active = label + if label: + self._touch(label) + + # -- release ----------------------------------------------------------- # + def evict(self, label: str, save: bool = True) -> bool: + """Drop a resident worker, persisting its database first.""" + client = self._clients.pop(label, None) + if client is None: + return False + if label in self._lru: + self._lru.remove(label) + if save: + try: # persist analysis + edits so the next open is a load + client.call("idb_save") + except Exception: # noqa: BLE001 -- evict regardless + pass + try: + client.close() + except Exception: # noqa: BLE001 + pass + self.evicted.append(label) + return True + + def _evictable(self, protect: str | None) -> str | None: + for label in self._lru: # least-recently-used first + if label == protect or label == self.active or label in self._pinned: + continue + return label + return None + + def _enforce_budget(self, protect: str | None = None) -> int: + """Evict LRU workers until the pool fits its budget. Returns how many.""" + n = 0 + while self.memory_mb() > self.budget_mb: + victim = self._evictable(protect) + if victim is None: # everything left is active/pinned/protected + break + self.evict(victim) + n += 1 + return n + + def close_all(self, save: bool = True) -> None: + for label in list(self._clients): + self.evict(label, save=save) + self.active = None + + # -- introspection ------------------------------------------------------ # + def status(self) -> list[dict]: + """Per-binary residency for the switcher UI.""" + out = [] + for ref in self.project.refs: + client = self._clients.get(ref.label) + out.append({ + "label": ref.label, + "source": ref.source, + "resident": client is not None, + "pinned": ref.label in self._pinned, + "active": ref.label == self.active, + "analysed": self.project.has_db(ref), + "memory_mb": self._mem(client) if client is not None else 0, + }) + return out + + def __repr__(self) -> str: # pragma: no cover - debug aid + return (f"<WorkerPool {len(self._clients)}/{len(self.project.refs)} resident " + f"{self.memory_mb()}/{self.budget_mb}MB active={self.active}>") diff --git a/idatui/worker_client.py b/idatui/worker_client.py index 27fa660..586ff21 100644 --- a/idatui/worker_client.py +++ b/idatui/worker_client.py @@ -117,7 +117,19 @@ class WorkerClient: time.sleep(0.2) raise IDAConnectionError("worker did not become ready in time") - def close(self) -> None: + @property + def pid(self) -> int | None: + """The worker process id (for memory accounting), or None if not spawned.""" + return self._proc.pid if self._proc is not None else None + + def close(self, grace: float = 20.0) -> None: + """Shut the worker down cleanly. + + After ``__shutdown__`` the worker still has to ``close_database()``, which + re-packs the ``.i64`` and removes the unpacked ``.id0/.id1/...`` scratch. + Signalling it before that finishes is what leaves databases wedged, so + wait out the grace period first and only escalate if it really is stuck. + """ with self._lock: s = self._sock self._sock = None @@ -132,13 +144,16 @@ class WorkerClient: pass if self._proc is not None: try: - self._proc.terminate() - self._proc.wait(timeout=5) - except Exception: # noqa: BLE001 + self._proc.wait(timeout=grace) # let it close the DB properly + except Exception: # noqa: BLE001 -- TimeoutExpired: it's stuck try: - self._proc.kill() + self._proc.terminate() + self._proc.wait(timeout=5) except Exception: # noqa: BLE001 - pass + try: + self._proc.kill() + except Exception: # noqa: BLE001 + pass # -- the call surface -------------------------------------------------- # def call(self, tool: str, *, timeout: float | None = None, **args) -> Any: diff --git a/tests/test_pool.py b/tests/test_pool.py new file mode 100644 index 0000000..d4347d9 --- /dev/null +++ b/tests/test_pool.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 +"""Unit tests for idatui.pool (worker residency: LRU + memory budget). + +Pure stdlib with a fake client injected, so the eviction policy is testable +without spawning real idalib workers. + + python tests/test_pool.py +""" +import os +import sys +import tempfile + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from idatui.pool import WorkerPool # noqa: E402 +from idatui.project import Project # noqa: E402 + +PASS = FAIL = 0 + + +def check(name, cond, detail=""): + global PASS, FAIL + if cond: + PASS += 1 + print(f" ok {name}") + else: + FAIL += 1 + print(f" FAIL {name} {detail}") + + +class FakeClient: + """Stands in for a WorkerClient: records saves/closes, reports fixed memory.""" + + def __init__(self, ref, mem=100): + self.ref = ref + self.mem = mem + self.saved = 0 + self.closed = False + self.connected = False + + def connect(self, progress=None, **kw): + self.connected = True + return self + + def call(self, tool, **kw): + if tool == "idb_save": + self.saved += 1 + return {} + + def close(self, grace=None): + self.closed = True + + +def _mkproject(tmp, n=4, size=64): + src = os.path.join(tmp, "src") + os.makedirs(src, exist_ok=True) + paths = [] + for i in range(n): + p = os.path.join(src, f"bin{i}") + with open(p, "wb") as f: + f.write(b"\x7fELF" + bytes(size)) + paths.append(p) + return Project.create(os.path.join(tmp, "p.json"), paths, name="p") + + +def main() -> int: + with tempfile.TemporaryDirectory() as tmp: + proj = _mkproject(tmp) + made = {} + + def spawn(ref, ttl): + c = FakeClient(ref, mem=100) + made[ref.label] = c + return c + + pool = WorkerPool(proj, budget_mb=350, spawn=spawn, + mem_fn=lambda c: c.mem) + + # -- lazy spawn + reuse -------------------------------------------- # + a = pool.get("bin0") + check("get() spawns a worker on first use", a is made["bin0"] and a.connected) + check("get() stages the binary first", + os.path.isfile(proj.by_label("bin0").staged)) + check("get() reuses the resident worker", pool.get("bin0") is a) + check("resident() reports it", pool.resident() == ["bin0"], pool.resident()) + + # -- LRU ordering ---------------------------------------------------- # + pool.get("bin1") + pool.get("bin2") + pool.get("bin0") # touch: bin0 becomes most-recent + check("LRU order tracks use", pool.resident() == ["bin1", "bin2", "bin0"], + pool.resident()) + + # -- budget eviction -------------------------------------------------- # + check("pool reports its memory", pool.memory_mb() == 300, pool.memory_mb()) + pool.get("bin3") # 400MB > 350MB budget -> evict LRU (bin1) + check("exceeding the budget evicts the least-recently-used", + pool.evicted == ["bin1"] and not pool.is_resident("bin1"), + f"evicted={pool.evicted} resident={pool.resident()}") + check("the just-spawned worker is never the victim", pool.is_resident("bin3")) + check("eviction saves the database first", made["bin1"].saved == 1) + check("eviction closes the worker", made["bin1"].closed) + check("pool is back within budget", pool.memory_mb() <= pool.budget_mb, + f"{pool.memory_mb()}/{pool.budget_mb}") + + # -- the active binary is never evicted ------------------------------- # + pool.set_active("bin2") + check("set_active touches the LRU", pool.resident()[-1] == "bin2", + pool.resident()) + pool.get("bin1") # over budget again -> must evict, but not bin2 + check("the active binary survives eviction", pool.is_resident("bin2"), + f"resident={pool.resident()}") + + # -- pinning ---------------------------------------------------------- # + pool.close_all() + pool2 = WorkerPool(proj, budget_mb=250, spawn=spawn, mem_fn=lambda c: c.mem) + pool2.get("bin0") + pool2.pin("bin0") + pool2.get("bin1") + pool2.get("bin2") # 300 > 250 -> evict, but bin0 is pinned + check("pinned binaries are never evicted", pool2.is_resident("bin0"), + f"resident={pool2.resident()} evicted={pool2.evicted}") + check("an unpinned one went instead", "bin1" in pool2.evicted, pool2.evicted) + + # -- everything pinned/active: stop evicting rather than thrash -------- # + pool2.pin("bin2") + pool2.set_active("bin2") + n_before = len(pool2.evicted) + pool2._enforce_budget() + check("nothing evictable -> gives up instead of thrashing", + len(pool2.evicted) == n_before, pool2.evicted) + + # -- status for the switcher UI ---------------------------------------- # + st = {s["label"]: s for s in pool2.status()} + check("status() covers every project binary", len(st) == 4, list(st)) + check("status() marks resident/pinned/active", + st["bin0"]["resident"] and st["bin0"]["pinned"] + and st["bin2"]["active"] and not st["bin3"]["resident"], + f"{st}") + + # -- teardown ----------------------------------------------------------- # + pool2.close_all() + check("close_all() closes every worker", + not pool2.resident() and all(c.closed for c in made.values())) + check("close_all() clears the active binary", pool2.active is None) + + # -- unknown label -------------------------------------------------------- # + try: + pool2.get("nope") + check("an unknown label raises KeyError", False, "no raise") + except KeyError: + check("an unknown label raises KeyError", True) + + # -- default budget comes from the project's memory_pct ------------------- # + pool3 = WorkerPool(proj, spawn=spawn, mem_fn=lambda c: c.mem) + check("default budget is derived, not a fixed worker count", + pool3.budget_mb >= 256, pool3.budget_mb) + + print(f"\n{PASS} passed, {FAIL} failed") + return 1 if FAIL else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) |