From e7ad2de0e7de8bbaac8e5ba6eaf8c698b1a4f862 Mon Sep 17 00:00:00 2001 From: Kim Morrison Date: Fri, 6 Nov 2026 23:07:06 +0000 Subject: [PATCH] fix: publish on-demand Lake cache downloads atomically --- src/lake/Lake/Config/Cache.lean | 31 ++++++-- tests/lake/tests/cacheDownload/Probe.lean | 11 +++ tests/lake/tests/cacheDownload/test.py | 87 +++++++++++++++++++++++ tests/lake/tests/cacheDownload/test.sh | 8 +++ 4 files changed, 130 insertions(+), 7 deletions(-) create mode 100644 tests/lake/tests/cacheDownload/Probe.lean create mode 100644 tests/lake/tests/cacheDownload/test.py create mode 100755 tests/lake/tests/cacheDownload/test.sh diff --git a/src/lake/Lake/Config/Cache.lean b/src/lake/Lake/Config/Cache.lean index e224b92cf9db..806faeaf7234 100644 --- a/src/lake/Lake/Config/Cache.lean +++ b/src/lake/Lake/Config/Cache.lean @@ -553,15 +553,32 @@ end CacheToolchain /-! ## Remote Cache Service -/ +partial def createArtifactDownloadDir (path : FilePath) : IO FilePath := do + let pid ← IO.Process.getPID + let nonce ← IO.monoNanosNow + let dir := path.addExtension s!"{pid}-{nonce}.tmp" + try + IO.FS.createDir dir + return dir + catch + | .alreadyExists .. => createArtifactDownloadDir path + | e => throw e + /-- **For internal use only.** -/ public def downloadArtifactCore (hash : Hash) (url : String) (path : FilePath) : LogIO Unit := do - download url path - let actualHash ← computeFileHash path - if actualHash != hash then - let errPos ← getLogPos - logError s!"{path}: downloaded artifact hash mismatch, got {actualHash}" - IO.FS.removeFile path - throw errPos + createParentDirs path + -- `download` removes its destination, so reserve a directory rather than a file. + let tmpDir ← createArtifactDownloadDir path + let tmp := tmpDir / "artifact" + try + download url tmp + let actualHash ← computeFileHash tmp + if actualHash != hash then + error s!"{path}: downloaded artifact hash mismatch, got {actualHash}" + IO.FS.rename tmp path + finally + removeFileIfExists tmp + IO.FS.removeDir tmpDir /-- Uploads a file to an online bucket using the Amazon S3 protocol. -/ def uploadS3 diff --git a/tests/lake/tests/cacheDownload/Probe.lean b/tests/lake/tests/cacheDownload/Probe.lean new file mode 100644 index 000000000000..e59a62502fc5 --- /dev/null +++ b/tests/lake/tests/cacheDownload/Probe.lean @@ -0,0 +1,11 @@ +import Lake.Config.Cache + +/-! Test on-demand cache download publication for #16301. -/ + +open System Lake + +def main (args : List String) : IO UInt32 := do + let [url, path, contents] := args + | throw <| IO.userError "expected URL PATH CONTENTS" + let result ← (downloadArtifactCore (Hash.ofByteArray contents.toUTF8) url path).toBaseIO + return if result.isSome then 0 else 1 diff --git a/tests/lake/tests/cacheDownload/test.py b/tests/lake/tests/cacheDownload/test.py new file mode 100644 index 000000000000..4a342624313e --- /dev/null +++ b/tests/lake/tests/cacheDownload/test.py @@ -0,0 +1,87 @@ +"""An on-demand download becomes visible only after completion and validation.""" + +import http.server +import os +import pathlib +import subprocess +import tempfile +import threading +import unittest + + +class Download(unittest.TestCase): + def test_successful_download(self): + self.download("success") + + def test_hash_mismatch(self): + self.download("corrupt") + + def test_interrupted_transfer(self): + self.download("incomplete") + + def download(self, outcome): + expected = b"complete artifact bytes" + payload = b"incorrect artifact bytes" if outcome == "corrupt" else expected + for existing in (False, True): + with self.subTest(existing=existing), tempfile.TemporaryDirectory() as temporary: + path = pathlib.Path(temporary) / "artifacts/artifact.art" + previous = b"previous artifact bytes" + if existing: + path.parent.mkdir() + path.write_bytes(previous) + started, release = threading.Event(), threading.Event() + + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload[:4]) + self.wfile.flush() + started.set() + release.wait(30) + if outcome != "incomplete": + self.wfile.write(payload[4:]) + + def log_message(self, *_): + pass + + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + server.daemon_threads = True + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + child = subprocess.Popen( + [os.environ.get("LEAN", "lean"), "-Dinterpreter.prefer_native=false", "--run", "Probe.lean", + f"http://127.0.0.1:{server.server_port}/artifact", str(path), expected.decode()], + stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, + ) + try: + self.assertTrue(started.wait(30), "native downloader did not reach the server") + if existing: + self.assertEqual(path.read_bytes(), previous, "download changed the final name before completion") + else: + self.assertFalse(path.exists(), "partial download is visible under the final name") + finally: + release.set() + try: + out, err = child.communicate(timeout=30) + except subprocess.TimeoutExpired: + child.kill() + child.communicate() + raise + finally: + server.shutdown() + server.server_close() + thread.join() + self.assertEqual(child.returncode == 0, outcome == "success", out + err) + if outcome == "success": + self.assertEqual(path.read_bytes(), expected) + elif existing: + self.assertEqual(path.read_bytes(), previous) + else: + self.assertFalse(path.exists()) + self.assertEqual(list(path.parent.iterdir()), [path] if path.exists() else [], "staging was not removed") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/lake/tests/cacheDownload/test.sh b/tests/lake/tests/cacheDownload/test.sh new file mode 100755 index 000000000000..1ff3e3d0eff6 --- /dev/null +++ b/tests/lake/tests/cacheDownload/test.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +source ../common.sh +PYTHON="${PYTHON:-python3}" +if ! "$PYTHON" -c 'import sys; sys.exit(sys.version_info < (3, 11))' 2> /dev/null; then + echo "SKIPPED: $PYTHON 3.11 or later not found" + exit 0 +fi +"$PYTHON" test.py