fix: persist the sandbox namespace with a single dill dump

synalinks/src/sandboxes/mirage_sandbox.py

@@ -746,20 +746,52 @@ def _stub(*args, **kwargs):
     else:
         exec(compile(tree, "<sandbox>", "exec"), ns)
 finally:
-    keep = {}
-    for key, item in list(ns.items()):
-        if key == "__builtins__":
-            continue
-        try:
-            dill.dumps(item)
-            keep[key] = item
-        except Exception:
-            pass
+    # Persist the namespace with ONE dill.dumps. Pickling item by item is
+    # quadratic: every sandbox-defined function is pickled together with a
+    # copy of its (shared) globals, so N functions cost N full-namespace
+    # pickles (25 functions ~3 s, 100 ~14 s, ~200 hits RecursionError). A
+    # single dump memoizes the shared globals once. Only when that fails do
+    # we fall back to filtering out the unpicklable items one by one.
+    keep = {k: v for k, v in ns.items() if k != "__builtins__"}
     try:
-        with open(state, "wb") as fh:
-            dill.dump(keep, fh)
-    except Exception as exc:
-        print("persist-warn: " + repr(exc), file=sys.stderr)
+        blob = dill.dumps(keep)
+    except Exception:
+        # Something in the namespace is unpicklable (an open file, a
+        # generator, ...). Drop those items, and pickle functions with
+        # ``recurse=True`` so they carry only the globals they reference
+        # rather than the whole namespace: otherwise every sandbox-defined
+        # function would be lost along with the offending item. Restored
+        # functions are re-homed onto the live namespace on the next run
+        # anyway, so the reduced globals are never observable.
+        keep = {}
+        for key, item in list(ns.items()):
+            if key == "__builtins__":
+                continue
+            try:
+                dill.dumps(item, recurse=True)
+                keep[key] = item
+            except Exception:
+                pass
+        try:
+            blob = dill.dumps(keep, recurse=True)
+        except Exception as exc:
+            blob = None
+            print("persist-warn: " + repr(exc), file=sys.stderr)
+    if blob is not None:
+        # Write to a sibling temp file and rename so a run killed mid-write
+        # (host timeout) or a concurrent run can never leave the state file
+        # truncated or half-written.
+        tmp = state + ".tmp." + str(os.getpid())
+        try:
+            with open(tmp, "wb") as fh:
+                fh.write(blob)
+            os.replace(tmp, state)
+        except Exception as exc:
+            print("persist-warn: " + repr(exc), file=sys.stderr)
+            try:
+                os.unlink(tmp)
+            except OSError:
+                pass
     result_path = config.get("result")
     if result_path:
         try:

synalinks/src/sandboxes/mirage_sandbox_test.py

@@ -97,6 +97,36 @@ async def test_functions_and_classes_persist(self):
         result = await sandbox.run("print(sq(4), math.floor(2.7), P.v)")
         self.assertIn("16 2 3", result.stdout)
 
+    async def test_many_functions_persist(self):
+        """Persisting must not be quadratic in the number of definitions.
+
+        dill pickles each function together with a copy of its globals, so
+        pickling the namespace item by item re-pickles every sibling function
+        once per function: 120 definitions took ~4 s and a few hundred hit
+        `RecursionError`, silently losing the whole namespace. One dump of the
+        namespace memoizes the shared globals instead."""
+        sandbox = MirageSandbox(timeout=_TIMEOUT)
+        await sandbox.run(
+            "\n".join(f"def f{i}(x):\n    return x + {i}" for i in range(120))
+        )
+        result = await sandbox.run("print(f0(1), f119(1))")
+        self.assertIn("1 120", result.stdout)
+
+    async def test_unpicklable_value_does_not_drop_the_namespace(self):
+        """One unpicklable object must not take the definitions with it.
+
+        A generator (an open socket, a live handle) cannot be pickled. The
+        fallback drops it and pickles the rest with `recurse=True` so each
+        function carries only the globals it references, instead of the
+        offending object along with the whole namespace."""
+        sandbox = MirageSandbox(timeout=_TIMEOUT)
+        await sandbox.run(
+            "gen = (i for i in range(3))\ndef survivor(x):\n    return x * 3"
+        )
+        result = await sandbox.run("print(survivor(5))")
+        self.assertIn("15", result.stdout)
+        self.assertIsNone(result.error)
+
     async def test_function_sees_names_defined_in_later_runs(self):
         """A restored function shares the live namespace, like a real REPL.