projects: worker pool with memory-budgeted residency (phase 1b)

idatui/pool.py — WorkerPool keeps one live worker per project binary:
* lazy spawn on first use; staging + a scratch sweep happen first, so a database
  wedged by a previously hard-killed worker reopens instead of crash-looping.
* residency is bounded by a MEMORY BUDGET (default project.memory_pct of RAM),
  not a worker count — a count is the wrong knob when one project holds a 50KB
  helper and a 6MB crypto lib. Cost is measured per worker from
  /proc/<pid>/smaps_rollup (PSS, which splits shared pages so summing means
  something).
* over budget -> evict least-recently-used, never the active or pinned binary,
  and give up rather than thrash when nothing is evictable. Eviction calls
  idb_save first, so returning to a binary is a DB load, not a re-analysis.
* status() feeds the switcher UI (resident/pinned/active/analysed/memory).

worker_client: fix the wedge-file bug this depends on. close() sent
__shutdown__ and then IMMEDIATELY SIGTERMed, killing the worker mid
close_database() — which is what leaves the unpacked .id0/.id1/... behind and
makes the .i64 refuse to reopen. Now it waits out a grace period (the worker
returns from serve() on __shutdown__ and closes the DB in its finally) and only
escalates if it's genuinely stuck. Also expose .pid for memory accounting.
Verified: a pilot run that used to leave echo.id0/.id1/.id2/.nam/.til now leaves
only echo.i64, and teardown is no slower (14 checks in 5s).

tests/test_pool.py: 22 checks with an injected fake client — LRU order, budget
eviction, active/pinned protection, save-before-close, thrash avoidance,
status(), teardown. Pure stdlib, no idalib.

2 files changed, 225 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: