Make the Lean service safe for concurrent checks of one project

docker/lean/README.md

@@ -12,8 +12,25 @@ the decision is [ADR 0003](../../docs/adr/0003-run-lean-as-a-pinned-service.md).
 | `Dockerfile` | toolchain, `lake exe cache get`, the REPL, the wrapper |
 | `server.py` | routes and request handling |
 | `repl.py` | warm REPL pool, deadlines, output cap |
+| `build.py` | compiles a project's own modules with `lean -o`; per-project locks, atomic oleans |
 | `normalize.py` | pure: REPL answers to replies, and the Lean commands sent |
 | `test_normalize.py` | `python3 -B -m unittest test_normalize` (no Lean needed) |
+| `test_build.py`, `test_server_concurrency.py` | build order, staging, and concurrent checks (no Lean needed) |
+
+## Concurrent checks
+
+The server is threaded, so a caller may send several `/check` requests for one
+project at once. Staging a project and compiling its `Supermath.*` modules run
+under a lock per project (keyed by the real path of its `lean/` directory), so
+concurrent requests for the same project build one at a time: a module two
+files import is compiled once, no `.olean` is written twice, and a restage
+never rewrites sources under a compile. Requests for different projects hold
+different locks and build in parallel. The lock covers staging and the build
+only; elaborating the file afterwards just reads the finished oleans and runs
+concurrently, bounded by `LEAN_POOL_SIZE`. Oleans are also written atomically
+(to a `.tmp` file renamed into place), so a failed or timed-out compile leaves
+no partial olean. The locks are per process; run one service per
+`LEAN_BUILDS_DIR` and projects volume.
 
 Bumping Mathlib: change `rev` in `lakefile.toml` and the toolchain in
 `lean-toolchain`, rebuild, commit the new `lake-manifest.json`, and re-check

docker/lean/build.py

@@ -6,13 +6,22 @@
 resolve Mathlib's dependencies. The packages are already built, so all a
 project needs is `lean -o` on its own modules, in dependency order, against a
 `LEAN_PATH` that lists the baked packages and the project's build directory.
+
+Concurrency: a project's build directory is not safe for two writers, so
+callers hold `project_lock(...)` across staging and `build_modules` for one
+project (the server does, in `build_package`). Under that lock requests for one
+project build one at a time and requests for different projects run in
+parallel. `build_modules` also writes each olean atomically, to a temporary
+file renamed into place once Lean succeeds, so a failed or timed-out compile
+never leaves a half-written olean that `stale` would take for fresh.
 """
 
 import glob
 import hashlib
 import os
 import shutil
 import subprocess
+import threading
 
 import normalize as N
 
@@ -27,6 +36,23 @@ def __init__(self, text):
         self.text = text
 
 
+_LOCKS = {}
+_LOCKS_GUARD = threading.Lock()
+
+
+def project_lock(source_dir):
+    """The lock that serialises staging and building the package at
+    `source_dir`, one per real path, created on first use.
+
+    Hold it from `stage` through `build_modules`: `stage` rewrites sources and
+    `build_modules` writes oleans in the same directory, and neither may run
+    while the other, or a second copy of itself, does for the same project.
+    """
+    key = os.path.realpath(source_dir)
+    with _LOCKS_GUARD:
+        return _LOCKS.setdefault(key, threading.Lock())
+
+
 def stage(source_dir, builds_root):
     """Mirrors the `.lean` files of the package at `source_dir` into a build
     directory of its own under `builds_root`, and returns that directory.
@@ -144,15 +170,23 @@ def build_modules(lean_dir, modules, packages, timeout, run=subprocess.run):
             deps = N.project_imports(f.read())
         if not stale(lean_dir, module, deps):
             continue
-        olean = module_olean(lean_dir, module)
+        olean, ilean = module_olean(lean_dir, module), module_olean(lean_dir, module, "ilean")
         os.makedirs(os.path.dirname(olean), exist_ok=True)
-        done = run(
-            ["lean", "-o", olean, "-i", module_olean(lean_dir, module, "ilean"),
-             module_source(lean_dir, module)],
-            cwd=lean_dir, env=env, capture_output=True, text=True, timeout=timeout,
-        )
-        if done.returncode != 0:
-            raise BuildError(f"{module} does not compile:\n{(done.stdout + done.stderr)[-1500:]}")
+        outputs = [(olean + ".tmp", olean), (ilean + ".tmp", ilean)]
+        try:
+            done = run(
+                ["lean", "-o", outputs[0][0], "-i", outputs[1][0], module_source(lean_dir, module)],
+                cwd=lean_dir, env=env, capture_output=True, text=True, timeout=timeout,
+            )
+            if done.returncode != 0:
+                raise BuildError(f"{module} does not compile:\n{(done.stdout + done.stderr)[-1500:]}")
+            for temporary, final in outputs:
+                if os.path.exists(temporary):
+                    os.replace(temporary, final)
+        finally:
+            for temporary, _ in outputs:
+                if os.path.exists(temporary):
+                    os.remove(temporary)
     return lean_path(paths)
 
 

docker/lean/server.py

@@ -18,6 +18,12 @@
 lie inside one of the directories SUPERMATH_ROOTS lists, or it is refused.
 Such a project is only ever read: its package is built in a copy under
 LEAN_BUILDS_DIR (`build.stage`), so no build output lands in the checkout.
+
+Concurrency: the server handles requests on threads. Staging and compiling a
+project's modules run under that project's lock (`build_package`), so
+concurrent checks of the same project build serially, never writing one olean
+twice or restaging under a compile, while checks of different projects run in
+parallel. Elaboration after the build only reads oleans and is not locked.
 """
 
 import json
@@ -106,6 +112,20 @@ def ensure_package(project):
     return lean_dir
 
 
+def build_package(project, mods, timeout, run=subprocess.run):
+    """Stages and builds `project`'s package; returns (lean_dir, LEAN_PATH).
+
+    Staging and the build run under the project's lock (`build.project_lock`),
+    so concurrent requests for one project build one at a time, and a later
+    staging cannot rewrite sources under a compile. Requests for other
+    projects take other locks and build in parallel. The lock is released
+    before the caller elaborates: that only reads the finished oleans.
+    """
+    with build.project_lock(os.path.join(project, "lean")):
+        lean_dir = ensure_package(project)
+        return lean_dir, build.build_modules(lean_dir, mods, PACKAGES, timeout, run)
+
+
 def run_guarded(operation):
     """Runs `operation()`; maps REPL failures to a service status."""
     started = time.monotonic()
@@ -142,9 +162,8 @@ def release():
             release()
             raise
         return repl, answer, release
-    lean_dir = ensure_package(project)
     try:
-        path = build.build_modules(lean_dir, mods, PACKAGES, timeout)
+        lean_dir, path = build_package(project, mods, timeout)
     except build.BuildError as failed:
         fail = {"messages": [{"severity": "error", "pos": {"line": 1, "column": 0},
                               "data": "an imported project module failed to build:\n" + failed.text}]}

docker/lean/test_server_concurrency.py

@@ -0,0 +1,165 @@
+"""Tests for the server's package build under concurrent requests.
+
+    python3 -I -m unittest discover -s docker/lean
+
+No Lean is needed: `build_package` takes the `run` that compiles a module, and
+the fakes here record who is compiling what, and when.
+"""
+
+import os
+import subprocess
+import tempfile
+import threading
+import time
+import unittest
+
+import build as B
+import server as S
+
+
+def read(path):
+    with open(path, encoding="utf-8") as f:
+        return f.read()
+
+
+def write(path, text):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(text)
+
+
+class Probe:
+    """A fake `lean -o` that notes how many compiles of one olean overlap."""
+
+    def __init__(self, pause=0.15):
+        self.pause, self.lock = pause, threading.Lock()
+        self.active, self.overlaps, self.compiles = {}, [], []
+
+    def __call__(self, argv, **kw):
+        olean = argv[2]
+        with self.lock:
+            self.active[olean] = self.active.get(olean, 0) + 1
+            self.compiles.append(olean)
+            if self.active[olean] > 1:
+                self.overlaps.append(olean)
+        time.sleep(self.pause)
+        write(olean, "olean")
+        write(argv[4], "ilean")
+        with self.lock:
+            self.active[olean] -= 1
+        return subprocess.CompletedProcess(argv, 0, "", "")
+
+
+def in_threads(jobs):
+    """Runs each job in its own thread; returns the exceptions they raised."""
+    errors = []
+
+    def guarded(job):
+        try:
+            job()
+        except BaseException as error:  # noqa: BLE001 - reported to the test
+            errors.append(error)
+
+    threads = [threading.Thread(target=guarded, args=(job,)) for job in jobs]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(30)
+    return errors
+
+
+class ConcurrentBuilds(unittest.TestCase):
+    def setUp(self):
+        tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(tmp.cleanup)
+        self.tmp = tmp.name
+        self.saved = (S.PROJECTS, S.BUILDS)
+        S.PROJECTS, S.BUILDS = os.path.join(self.tmp, "projects"), os.path.join(self.tmp, "builds")
+        self.addCleanup(lambda: setattr(S, "PROJECTS", self.saved[0]) or setattr(S, "BUILDS", self.saved[1]))
+
+    def project(self, name, volume):
+        base = S.PROJECTS if volume else os.path.join(self.tmp, "checkouts")
+        project = os.path.join(base, name)
+        write(os.path.join(project, "lean", "Supermath", "Shared.lean"), "import Mathlib\n")
+        write(os.path.join(project, "lean", "Supermath", "Top.lean"), "import Supermath.Shared\n")
+        return project
+
+    def test_two_checks_of_one_projects_volume_never_compile_a_module_at_once(self):
+        self.same_project_serialises(volume=True)
+
+    def test_two_checks_of_one_staged_project_never_compile_a_module_at_once(self):
+        self.same_project_serialises(volume=False)
+
+    def same_project_serialises(self, volume):
+        project = self.project("p", volume)
+        probe = Probe()
+        errors = in_threads([lambda: S.build_package(project, ["Supermath.Top"], 10, probe)] * 4)
+        self.assertEqual(errors, [])
+        self.assertEqual(probe.overlaps, [], "one olean was written by two compiles at once")
+        self.assertEqual(len(probe.compiles), 2, "the later requests reuse what the first built")
+
+    def test_a_restage_waits_for_the_build_in_progress(self):
+        project = self.project("p", volume=False)
+        source = os.path.join(project, "lean", "Supermath", "Shared.lean")
+        seen = []
+
+        def run(argv, **kw):
+            before = read(argv[-1])
+            time.sleep(0.15)
+            seen.append(before == read(argv[-1]))
+            return Probe(0)(argv, **kw)
+
+        def edit_then_build():
+            time.sleep(0.05)
+            write(source, "import Mathlib\n-- edited\n")
+            os.utime(source, (2_000_000_000, 2_000_000_000))
+            S.build_package(project, ["Supermath.Top"], 10, run)
+
+        errors = in_threads([lambda: S.build_package(project, ["Supermath.Top"], 10, run), edit_then_build])
+        self.assertEqual(errors, [])
+        self.assertTrue(seen and all(seen), "a source changed under a compile")
+
+    def test_checks_of_different_projects_build_at_the_same_time(self):
+        first, second = self.project("a", True), self.project("b", False)
+        both = threading.Barrier(2, timeout=5)
+
+        def run(argv, **kw):
+            both.wait()  # passes only while the other project is compiling too
+            return Probe(0)(argv, **kw)
+
+        errors = in_threads([
+            lambda: S.build_package(first, ["Supermath.Shared"], 10, run),
+            lambda: S.build_package(second, ["Supermath.Shared"], 10, run),
+        ])
+        self.assertEqual(errors, [], "different projects were serialised")
+
+
+class AtomicOleans(unittest.TestCase):
+    def test_a_failed_compile_leaves_no_olean_behind(self):
+        with tempfile.TemporaryDirectory() as lean_dir:
+            write(B.module_source(lean_dir, "Supermath.A"), "import Mathlib\n")
+
+            def run(argv, **kw):
+                write(argv[2], "half an olean")
+                return subprocess.CompletedProcess(argv, 1, "error", "")
+
+            with self.assertRaises(B.BuildError):
+                B.build_modules(lean_dir, ["Supermath.A"], [], 10, run)
+            self.assertFalse(os.path.exists(B.module_olean(lean_dir, "Supermath.A")))
+            self.assertEqual(os.listdir(os.path.dirname(B.module_olean(lean_dir, "Supermath.A"))), [])
+
+    def test_a_timed_out_compile_leaves_no_olean_behind(self):
+        with tempfile.TemporaryDirectory() as lean_dir:
+            write(B.module_source(lean_dir, "Supermath.A"), "import Mathlib\n")
+
+            def run(argv, **kw):
+                write(argv[2], "half an olean")
+                raise subprocess.TimeoutExpired(argv, 1)
+
+            with self.assertRaises(subprocess.TimeoutExpired):
+                B.build_modules(lean_dir, ["Supermath.A"], [], 10, run)
+            self.assertEqual(os.listdir(os.path.dirname(B.module_olean(lean_dir, "Supermath.A"))), [])
+
+
+if __name__ == "__main__":
+    unittest.main()