fix: publish Lake cache artifacts and metadata atomically

src/lake/Lake/Build/Common.lean

@@ -480,7 +480,7 @@ public def Cache.saveArtifact
       IO.setAccessRights file ⟨r, r, r⟩
       unless (← cacheFile.pathExists) do
         createParentDirs cacheFile
-        atomicWrite cacheFile normalized.toByteArray
+        writeFileAtomic cacheFile normalized (overwrite := false) (readOnly := true)
       IO.setAccessRights cacheFile ⟨r, r, r⟩
       writeFileHash file hash
       let mtime ← getMTime cacheFile
@@ -500,37 +500,14 @@ public def Cache.saveArtifact
         if let .error e ← (IO.FS.hardLink file cacheFile).toBaseIO then
           -- other functions can race to create the file
           unless e matches .alreadyExists .. do
-            atomicWrite cacheFile contents
+            writeBinFileAtomic cacheFile contents (overwrite := false) (readOnly := true) (exe := exe)
       IO.setAccessRights cacheFile ⟨r, r, r⟩
       writeFileHash file hash
       let mtime ← getMTime cacheFile
       let path := if useLocalFile then file else cacheFile
       return {descr, name := file.toString, path, mtime}
   catch e =>
     error s!"failed to cache artifact: {e}"
-where atomicWrite cacheFile contents := do
-  let tmpFile := cacheFile.addExtension "copy.tmp"
-  try discard <| IO.FS.Handle.mk tmpFile .writeNew catch _ => pure ()
-  let h ← IO.FS.Handle.mk tmpFile .readWrite
-  h.lock
-  if (← cacheFile.pathExists) then
-    try IO.FS.removeFile tmpFile catch _ => pure ()
-  else
-    h.rewind
-    -- A racing rename may leave us holding a handle to the cached artifact.
-    -- Thus, we must not write if the file already contains the expected contents.
-    unless (← h.readBinToEnd.toBaseIO).toOption.any (· == contents) do
-      h.rewind
-      -- Nonatomic write under the lock.
-      h.truncate
-      h.write contents
-      h.flush
-    -- Atomic rename. Must not overwrite because the `tmpFile` could be
-    -- an empty temporary file created by a racing rename and then a racing write.
-    try IO.FS.renameIfNew tmpFile cacheFile catch e =>
-      unless (← cacheFile.pathExists) do
-        throw e
-    try IO.FS.removeFile tmpFile catch _ => pure ()
 
 @[inline, inherit_doc Cache.saveArtifact]
 public def cacheArtifact

src/lake/Lake/CLI/Main.lean

@@ -566,9 +566,9 @@ protected def get : CliM PUnit := do
         return ws.defaultCacheService
     let map ← CacheMap.load file
     let localScope := pkg?.elim ws.root.cacheScope (·.cacheScope)
-    cache.writeMap localScope map service.name? (some remoteScope) overwrite
     let descrs ← map.collectOutputDescrs
     service.downloadArtifacts descrs cache remoteScope opts.forceDownload
+    cache.writeMap localScope map service.name? (some remoteScope) overwrite
   else
     let platform := opts.platform?.getD .system
     let toolchain := opts.toolchain?.getD ws.cacheToolchain
@@ -655,10 +655,10 @@ where
               if n = 0 || revs.size < n then "for any revision" else s!"in {n} revisions from HEAD"
             error s!"{remoteScope}: no outputs found {revisions}"
         return map
-    cache.writeMap pkg.cacheScope map service.name? (some remoteScope) overwrite
     unless opts.mappingsOnly do
       let descrs ← map.collectOutputDescrs
       service.downloadArtifacts descrs cache remoteScope opts.forceDownload
+    cache.writeMap pkg.cacheScope map service.name? (some remoteScope) overwrite
 
 def computeUploadService
   (service? : Option String) (lakeEnv : Env) (lakeCfg : LoadedLakeConfig)
@@ -831,13 +831,11 @@ protected def unstage : CliM PUnit := do
   let ok ← descrs.foldlM (init := true) fun ok descr => do
     let cachePath := artDir/ descr.relPath
     let stagingPath := stagingDir / descr.relPath
-    if (← cachePath.pathExists) then
-      if overwrite then
-        -- Cache artifacts are read-only, so the old artifact must be deleted first.
-        IO.FS.removeFile cachePath
-      else
-        return ok
-    match (← copyFile stagingPath cachePath |>.toBaseIO) with
+    if (← cachePath.pathExists) && !overwrite then return ok
+    let publish := do
+      let contents ← IO.FS.readBinFile stagingPath
+      writeBinFileAtomic cachePath contents overwrite (readOnly := true)
+    match (← publish |>.toBaseIO) with
     | .ok _ =>
       return ok
     | .error (.noFileOrDirectory ..) =>

src/lake/Lake/Config/Cache.lean

@@ -445,10 +445,7 @@ def writeOutputsCore
   createParentDirs file
   let out := {service?, scope? := remoteScope?, data := out : CacheOutput}
   let contents := (toJson out).pretty
-  if overwrite then
-    IO.FS.writeFile file contents
-  else
-    writeFileIfNew file contents
+  writeFileAtomic file contents overwrite
 
 /-- Cache the outputs corresponding to the given input for the package.  -/
 @[inline] public def writeOutputs
@@ -554,14 +551,19 @@ end CacheToolchain
 /-! ## Remote Cache Service -/
 
 /-- **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
+def downloadArtifactCoreWith (hash : Hash) (url : String) (path : FilePath) (overwrite : Bool) : LogIO Unit :=
+  withTempFileFor path fun tmp => do
+    download url tmp
+    let actualHash ← computeFileHash tmp
+    if actualHash != hash then
+      error s!"{tmp}: downloaded artifact hash mismatch, got {actualHash}"
+    let r := {read := true, write := false, execution := false}
+    IO.setAccessRights tmp ⟨r, r, r⟩
+    if overwrite then replaceFile tmp path else publishFileIfNew tmp path
+
+/-- **For internal use only.** -/
+public def downloadArtifactCore (hash : Hash) (url : String) (path : FilePath) : LogIO Unit :=
+  downloadArtifactCoreWith hash url path false
 
 /-- Uploads a file to an online bucket using the Amazon S3 protocol. -/
 def uploadS3
@@ -686,7 +688,7 @@ public def downloadArtifact
   logInfo s!"{scope}: downloading artifact {descr.hash}\
     \n  local path: {path}\
     \n  remote URL: {url}"
-  downloadArtifactCore descr.hash url path
+  downloadArtifactCoreWith descr.hash url path force
 
 public def uploadArtifact
   (contentHash : Hash) (art : FilePath) (service : CacheService) (scope : CacheServiceScope)
@@ -711,6 +713,7 @@ structure TransferInfo where
   path : FilePath
   /-- Additional paths for a downloaded artifact. -/
   extraPaths : Array FilePath := #[]
+  temporaryPath : FilePath := ""
 
 @[inline] def TransferInfo.addPath (self : TransferInfo) (path : FilePath) (extra := true) : TransferInfo :=
   if extra then
@@ -758,9 +761,6 @@ structure TransferState where
   -/
   settled : Std.HashSet Nat := {}
 
-@[inline] def tmpPath (path : FilePath) : FilePath :=
-  path.addExtension "tmp"
-
 /--
 The number of times a batch transfer is attempted: one initial attempt plus
 retries of the transfers that failed in ways a retry may fix (e.g., a dropped
@@ -844,27 +844,30 @@ where
       logError msg
       modify fun s => {s with settled := s.settled.insert i, didError := true}
   handleTransfer i info code out line : StateT TransferState LoggerIO Unit := do
-    let {url, hash, path, extraPaths} := info
+    let {url, hash, path, extraPaths, ..} := info
     match cfg.kind with
     | .get =>
       let logDownload :=
         logInfo s!"{cfg.scope}: downloaded artifact {hash}\
           \n  local path: {path}\
           \n  remote URL: {url}"
-      let tmpPath := tmpPath path
+      let tmpPath := info.temporaryPath
       match (← IO.FS.readBinFile tmpPath |>.toBaseIO) with
       | .ok contents =>
         logDownload
         let actualHash := Hash.ofByteArray contents
         if actualHash != hash then
           recordFailure i true s!"{tmpPath}: downloaded artifact hash mismatch, got {actualHash}"
-        else if let .error e ← IO.FS.rename tmpPath path |>.toBaseIO then
+        else if let .error e ← (do
+          let rights : IO.AccessRight := {read := true, write := false, execution := false}
+          IO.setAccessRights tmpPath ⟨rights, rights, rights⟩
+          publishFileIfNew tmpPath path) |>.toBaseIO then
           recordFailure i false s!"{path}: failed to persist temporary artifact: {e}"
         else
           unless extraPaths.isEmpty do
             -- Note: No intra-cache hard links (breaks permissions/pruning), so we copy
             for extraPath in extraPaths do
-              if let .error e ← IO.FS.writeBinFile extraPath contents |>.toBaseIO then
+              if let .error e ← writeBinFileAtomic extraPath contents (overwrite := false) (readOnly := true) |>.toBaseIO then
                 logError s!"{extraPath}: failed to copy artifact: {e}"
                 modify ({· with didError := true})
           recordSuccess i
@@ -890,7 +893,7 @@ where
       \n  remote URL: {info.url}"
     match cfg.kind with
     | .get =>
-      let tmpPath := tmpPath info.path
+      let tmpPath := info.temporaryPath
       unless code? matches .ok 404 do -- ignore response bodies on 404s
         if let .ok size := out.getAs Nat "size_download" then
           if size > 0 then
@@ -914,8 +917,7 @@ def runTransferAttempt
     match cfg.kind with
     | .get =>
       cfg.infos.forM fun info => do
-        let tmpPath := tmpPath info.path
-        removeFileIfExists tmpPath
+        let tmpPath := info.temporaryPath
         h.putStrLn s!"url = {info.url.quote}"
         h.putStrLn s!"-o {tmpPath.toString.quote}"
       h.flush
@@ -953,7 +955,7 @@ def runTransferAttempt
     logVerbose msg
   return s
 
-def transferArtifacts (cfg : TransferConfig) : LoggerIO Unit := do
+def transferArtifactsCore (cfg : TransferConfig) : LoggerIO Unit := do
   let action := match cfg.kind with | .get => "download" | .put => "upload"
   let mut didError := false
   let mut numSuccesses := 0
@@ -978,6 +980,20 @@ def transferArtifacts (cfg : TransferConfig) : LoggerIO Unit := do
   if didError then
     failure
 
+def transferArtifacts (cfg : TransferConfig) : LoggerIO Unit := do
+  let temporaryPaths ← IO.mkRef (#[] : Array FilePath)
+  try
+    let infos ← cfg.infos.mapM fun info => do
+      if cfg.kind matches .get then
+        let temporaryPath ← createTempFileFor info.path
+        temporaryPaths.modify (·.push temporaryPath)
+        return {info with temporaryPath}
+      else return info
+    transferArtifactsCore {cfg with infos}
+  finally
+    for path in (← temporaryPaths.get) do
+      removeFileIfExists path
+
 def reservoirArtifactsUrl (service : CacheService) (scope : CacheServiceScope) : String :=
   let endpoint :=
     match scope.impl with
@@ -1026,7 +1042,7 @@ where
     -- Note: No intra-cache hard links (breaks permissions/pruning), so we copy
     let contents ← IO.FS.readBinFile path
     for extraPath in extraPaths do
-      IO.FS.writeBinFile extraPath contents
+      writeBinFileAtomic extraPath contents (overwrite := false) (readOnly := true)
   fetchUrls url infos := IO.FS.withTempFile fun h path => do
     let body := Json.arr <| infos.map (toJson ·.hash)
     h.putStr body.compress
@@ -1161,7 +1177,7 @@ public def downloadRevisionOutputs?
   let some contents := contents?
     | return none
   createParentDirs path
-  IO.FS.writeFile path contents
+  writeFileAtomic path contents
   CacheMap.load path platform.isNone
 
 public def uploadRevisionOutputs

src/lake/Lake/Util/IO.lean

@@ -7,6 +7,7 @@ module
 
 prelude
 public import Init.System.IO
+import Init.System.Platform
 
 open System
 
@@ -24,6 +25,59 @@ public def removeFileIfExists (path : FilePath) : IO Unit := do
     | .noFileOrDirectory .. => pure ()
     | e => throw e
 
+/-- Reserve a unique temporary pathname beside {lean}`path`. The caller owns its removal. -/
+public partial def createTempFileFor (path : FilePath) : IO FilePath := do
+  createParentDirs path
+  let pid ← IO.Process.getPID
+  let nonce ← IO.monoNanosNow
+  let tmp := path.parent.getD "." / (".lake-cache-" ++ toString pid ++ "-" ++ toString nonce ++ ".tmp")
+  try
+    discard <| IO.FS.Handle.mk tmp .writeNew
+    return tmp
+  catch
+    | .alreadyExists .. => createTempFileFor path
+    | e => throw e
+
+/-- Run an action with a unique temporary file on the destination filesystem. -/
+public def withTempFileFor
+  [Monad m] [MonadFinally m] [MonadLiftT IO m]
+  (path : FilePath) (f : FilePath → m α)
+: m α := do
+  let tmp ← createTempFileFor path
+  try f tmp finally removeFileIfExists tmp
+
+/-- Publish a completed file, keeping any winner already at the destination. -/
+public def publishFileIfNew (tmp path : FilePath) : IO Unit := do
+  try IO.FS.hardLink tmp path catch
+    | .alreadyExists .. => pure ()
+    | _ =>
+      try IO.FS.renameIfNew tmp path catch e =>
+        unless (← path.pathExists) do throw e
+
+/-- Publish a completed temporary file, replacing the destination. -/
+public def replaceFile (tmp path : FilePath) : IO Unit := do
+  try IO.FS.rename tmp path catch e =>
+    if System.Platform.isWindows then
+      -- Windows cannot rename over a read-only cache entry. A missing entry is recoverable.
+      removeFileIfExists path
+      IO.FS.rename tmp path
+    else throw e
+
+/-- Publish complete bytes, or preserve an existing winner. Replacement is atomic on POSIX. -/
+public def writeBinFileAtomic
+  (path : FilePath) (contents : ByteArray) (overwrite := true) (readOnly exe := false)
+: IO Unit :=
+  withTempFileFor path fun tmp => do
+    IO.FS.writeBinFile tmp contents
+    if readOnly then
+      let rights := {read := true, write := false, execution := exe}
+      IO.setAccessRights tmp ⟨rights, rights, rights⟩
+    if overwrite then replaceFile tmp path else publishFileIfNew tmp path
+
+/-- Atomically replace a UTF-8 file, or preserve an existing winner. -/
+public def writeFileAtomic (path : FilePath) (contents : String) (overwrite := true) (readOnly := false) : IO Unit :=
+  writeBinFileAtomic path contents.toUTF8 overwrite readOnly
+
 /--
 Write the UTF-8 encoded string {lean}`content` to {lean}`path`.
 If the file already exists, does nothing.

tests/lake/tests/cacheAtomicWrite/lakefile.lean

@@ -2,8 +2,7 @@ import Lake
 open System Lake DSL
 
 /-!
-Test how the Lake artifact cache recovers from
-a temporary file left by an interrupted write.
+Test that publication ignores temporary files owned by another publisher.
 -/
 
 package test where

tests/lake/tests/cacheAtomicWrite/test.sh

@@ -3,8 +3,7 @@ source ../common.sh
 
 ./clean.sh
 
-# Test that caching a text artifact recovers from a temporary file
-# left behind by an interrupted write to the Lake artifact cache
+# Test that caching never reads or removes another publisher's temporary file
 
 TEST_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)"
 CACHE_DIR="$(norm_path "$TEST_DIR")/.lake/cache"
@@ -15,32 +14,24 @@ test_run build text
 art=$("$LAKE" query text)
 test_exp "$(norm_dirname "$art")" = "$CACHE_DIR/artifacts"
 test_cmd diff -u fixtures/source.txt "$art"
-tmp="$art.copy.tmp"
+tmp="$CACHE_DIR/artifacts/.lake-cache-$$-1.tmp"
 
-# Plants the given file as the temporary file of a fresh save,
-# re-caches the artifact, and verifies the result
+# Leave a different owner's staging file in place while re-caching the artifact.
 check_recovery() {
   rm -rf .lake/build "$CACHE_DIR/outputs"
   rm -f "$art"
   cat "$1" > "$tmp"
   test_run build text
   test_cmd diff -u fixtures/source.txt "$art"
-  test_exp ! -e "$tmp"
+  test_cmd diff -u "$1" "$tmp"
+  rm -f "$tmp"
 }
 
-echo "# TEST: complete temporary file"
+echo "# TEST: another publisher's temporary file"
 check_recovery fixtures/source.txt
-echo "# TEST: empty temporary file"
-check_recovery fixtures/empty.txt
-echo "# TEST: truncated temporary file"
-check_recovery fixtures/truncated.txt
-echo "# TEST: temporary file cut mid-character"
-check_recovery fixtures/midchar.txt
-echo "# TEST: oversized temporary file"
-check_recovery fixtures/oversized.txt
-
-# Verify no temporary files are left in the cache
-check_diff /dev/null <(ls -1 "$CACHE_DIR"/artifacts/*.tmp 2>/dev/null)
+
+# Include hidden staging files in the completed-invocation check.
+check_diff /dev/null <(find "$CACHE_DIR/artifacts" -name '.lake-cache-*.tmp' -print)
 
 # Cleanup
 rm -f produced.out