From 5b786f165d314f3b54056de85989e09e583d69d6 Mon Sep 17 00:00:00 2001
From: Wolf Vollprecht <w.vollprecht@gmail.com>
Date: Tue, 18 Feb 2025 14:01:09 +0100
Subject: [PATCH] feat: write prefix record atomically (#1063)

---
 crates/rattler_conda_types/Cargo.toml         |  2 +-
 .../rattler_conda_types/src/prefix_record.rs  | 43 +++++++++++++++++--
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/crates/rattler_conda_types/Cargo.toml b/crates/rattler_conda_types/Cargo.toml
index fafc2821d..edd4d76ae 100644
--- a/crates/rattler_conda_types/Cargo.toml
+++ b/crates/rattler_conda_types/Cargo.toml
@@ -36,6 +36,7 @@ serde-untagged = { workspace = true }
 serde_yaml = { workspace = true }
 smallvec = { workspace = true, features = ["serde", "const_new", "const_generics", "union"] }
 strum = { workspace = true, features = ["derive"] }
+tempfile = { workspace = true }
 thiserror = { workspace = true }
 tracing = { workspace = true }
 typed-path = { workspace = true }
@@ -50,7 +51,6 @@ fs-err = { workspace = true }
 rand = { workspace = true }
 insta = { workspace = true, features = ["yaml", "redactions", "toml", "glob", "filters"] }
 rattler_package_streaming = { path = "../rattler_package_streaming", default-features = false, features = ["rustls-tls"] }
-tempfile = { workspace = true }
 rstest = { workspace = true }
 assert_matches = { workspace = true }
 hex-literal = { workspace = true }
diff --git a/crates/rattler_conda_types/src/prefix_record.rs b/crates/rattler_conda_types/src/prefix_record.rs
index f722785c1..0fd32d822 100644
--- a/crates/rattler_conda_types/src/prefix_record.rs
+++ b/crates/rattler_conda_types/src/prefix_record.rs
@@ -4,7 +4,6 @@ use crate::package::FileMode;
 use crate::repo_data::RecordFromPath;
 use crate::repo_data_record::RepoDataRecord;
 use crate::PackageRecord;
-use fs_err::File;
 use rattler_digest::serde::SerializableHash;
 use serde::{Deserialize, Serialize};
 use serde_repr::{Deserialize_repr, Serialize_repr};
@@ -12,6 +11,7 @@ use serde_with::serde_as;
 use std::io::{BufWriter, Read};
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
+use tempfile::NamedTempFile;
 
 #[cfg(feature = "rayon")]
 use rayon::prelude::*;
@@ -234,8 +234,45 @@ impl PrefixRecord {
         path: impl AsRef<Path>,
         pretty: bool,
     ) -> Result<(), std::io::Error> {
-        let file = File::create(path.as_ref())?;
-        self.write_to(BufWriter::with_capacity(50 * 1024, file), pretty)
+        let path = path.as_ref();
+        let parent = path.parent().ok_or_else(|| {
+            std::io::Error::new(
+                std::io::ErrorKind::Other,
+                format!(
+                    "Failed to get parent directory of path '{}'",
+                    path.display()
+                ),
+            )
+        })?;
+
+        // Use a temporary file in the same directory for atomic writes
+        let temp_file = NamedTempFile::with_prefix_in("prefix_record_", parent).map_err(|e| {
+            std::io::Error::new(
+                std::io::ErrorKind::Other,
+                format!(
+                    "Failed to create temporary file in '{}': {}",
+                    parent.display(),
+                    e
+                ),
+            )
+        })?;
+
+        // Write to temp file with buffered writer
+        let writer = BufWriter::with_capacity(64 * 1024, &temp_file);
+        self.write_to(writer, pretty)?;
+
+        // Make sure that all data is written to disk
+        temp_file.as_file().sync_all()?;
+
+        // Atomically rename the temp file to the target path
+        temp_file.persist(path).map_err(|e| {
+            std::io::Error::new(
+                std::io::ErrorKind::Other,
+                format!("Failed to persist file {}: {}", path.display(), e),
+            )
+        })?;
+
+        Ok(())
     }
 
     /// Writes the contents of this instance to the file at the specified location.