1use std::fs::{File, OpenOptions};
23use std::io::{self, BufWriter, Write};
24use std::path::{Path, PathBuf};
25
26#[cfg(test)]
27thread_local! {
28 static TEST_FAILURE_STAGE: std::cell::Cell<Option<&'static str>> = const { std::cell::Cell::new(None) };
29}
30
31#[cfg(test)]
32fn fail_test_stage(stage: &'static str) -> io::Result<()> {
33 if TEST_FAILURE_STAGE.with(|value| value.get()) == Some(stage) {
34 return Err(io::Error::other(format!("injected {stage} failure")));
35 }
36 Ok(())
37}
38
39#[cfg(not(test))]
40#[inline]
41fn fail_test_stage(_stage: &'static str) -> io::Result<()> {
42 Ok(())
43}
44
45#[derive(Clone, Copy, Debug, Eq, PartialEq)]
47pub enum AtomicWriteDurability {
48 Namespace,
50 Flush,
54}
55
56#[derive(Clone, Copy, Debug, Eq, PartialEq)]
58pub struct AtomicWriteReceipt {
59 pub file_synced: bool,
61 pub namespace_synced: bool,
63}
64
65pub fn atomic_write(path: &Path, bytes: &[u8]) -> io::Result<()> {
67 atomic_write_with(path, |writer| writer.write_all(bytes))
68}
69
70pub fn atomic_write_with_mode(path: &Path, bytes: &[u8], mode: u32) -> io::Result<()> {
79 atomic_write_stream_with_durability_and_mode(
80 path,
81 AtomicWriteDurability::Flush,
82 Some(mode),
83 |writer| writer.write_all(bytes),
84 )
85 .map(|_| ())
86}
87
88pub fn atomic_copy_with_mode(
92 source: &Path,
93 destination: &Path,
94 mode: u32,
95) -> io::Result<AtomicWriteReceipt> {
96 let mut source = File::open(source)?;
97 atomic_write_stream_with_durability_and_mode(
98 destination,
99 AtomicWriteDurability::Flush,
100 Some(mode),
101 |writer| io::copy(&mut source, writer).map(|_| ()),
102 )
103}
104
105pub fn atomic_write_with_durability(
107 path: &Path,
108 bytes: &[u8],
109 durability: AtomicWriteDurability,
110) -> io::Result<AtomicWriteReceipt> {
111 atomic_write_stream_with_durability_and_mode(path, durability, None, |writer| {
112 writer.write_all(bytes)
113 })
114}
115
116pub(crate) fn atomic_write_with_durability_unlocked(
117 path: &Path,
118 bytes: &[u8],
119 durability: AtomicWriteDurability,
120) -> io::Result<AtomicWriteReceipt> {
121 atomic_write_stream_with_durability_and_mode_unlocked(path, durability, None, |writer| {
122 writer.write_all(bytes)
123 })
124}
125
126pub fn atomic_write_with<F>(path: &Path, write_fn: F) -> io::Result<()>
134where
135 F: FnOnce(&mut BufWriter<File>) -> io::Result<()>,
136{
137 atomic_write_stream_with_durability_and_mode(path, AtomicWriteDurability::Flush, None, write_fn)
138 .map(|_| ())
139}
140
141fn atomic_write_stream_with_durability_and_mode<F>(
142 path: &Path,
143 durability: AtomicWriteDurability,
144 mode: Option<u32>,
145 write_fn: F,
146) -> io::Result<AtomicWriteReceipt>
147where
148 F: FnOnce(&mut BufWriter<File>) -> io::Result<()>,
149{
150 #[cfg(windows)]
154 let _lock = crate::conditional_replace::acquire_lock(path)?;
155
156 atomic_write_stream_with_durability_and_mode_unlocked(path, durability, mode, write_fn)
157}
158
159fn atomic_write_stream_with_durability_and_mode_unlocked<F>(
160 path: &Path,
161 durability: AtomicWriteDurability,
162 mode: Option<u32>,
163 write_fn: F,
164) -> io::Result<AtomicWriteReceipt>
165where
166 F: FnOnce(&mut BufWriter<File>) -> io::Result<()>,
167{
168 let mut tmp = TempFile::create(path, mode)?;
169 let result = write_and_finalize(&mut tmp, durability, write_fn);
170 if let Err(err) = result {
171 let _ = std::fs::remove_file(&tmp.path);
172 return Err(err);
173 }
174 if let Err(err) = fail_test_stage("replace") {
175 let _ = std::fs::remove_file(&tmp.path);
176 return Err(err);
177 }
178 let replace_synced = match replace_temp_file(&tmp.path, path, durability) {
179 Ok(synced) => synced,
180 Err(err) => {
181 let _ = std::fs::remove_file(&tmp.path);
182 return Err(err);
183 }
184 };
185 let namespace_synced = match durability {
186 AtomicWriteDurability::Namespace => false,
187 AtomicWriteDurability::Flush => replace_synced || sync_parent_dir(path),
188 };
189 Ok(AtomicWriteReceipt {
190 file_synced: durability == AtomicWriteDurability::Flush,
191 namespace_synced,
192 })
193}
194
195fn write_and_finalize<F>(
196 tmp: &mut TempFile,
197 durability: AtomicWriteDurability,
198 write_fn: F,
199) -> io::Result<()>
200where
201 F: FnOnce(&mut BufWriter<File>) -> io::Result<()>,
202{
203 let file = tmp
204 .file
205 .take()
206 .ok_or_else(|| io::Error::other("atomic_io: temporary file handle was already consumed"))?;
207 let mut buf = BufWriter::new(file);
208 write_fn(&mut buf)?;
209 fail_test_stage("flush")?;
210 buf.flush()?;
211 let inner = buf.into_inner().map_err(|err| err.into_error())?;
212 if durability == AtomicWriteDurability::Flush {
213 inner.sync_all()?;
214 }
215 Ok(())
216}
217
218#[cfg(not(windows))]
219fn replace_temp_file(
220 temp: &Path,
221 destination: &Path,
222 _durability: AtomicWriteDurability,
223) -> io::Result<bool> {
224 std::fs::rename(temp, destination)?;
225 Ok(false)
226}
227
228#[cfg(windows)]
229fn replace_temp_file(
230 temp: &Path,
231 destination: &Path,
232 durability: AtomicWriteDurability,
233) -> io::Result<bool> {
234 use windows_sys::Win32::Storage::FileSystem::{
235 MoveFileExW, MOVEFILE_REPLACE_EXISTING, MOVEFILE_WRITE_THROUGH,
236 };
237
238 let temp_wide = crate::windows_path::wide_maybe_verbatim(temp);
239 let destination_wide = crate::windows_path::wide_maybe_verbatim(destination);
240 let mut flags = MOVEFILE_REPLACE_EXISTING;
241 if durability == AtomicWriteDurability::Flush {
242 flags |= MOVEFILE_WRITE_THROUGH;
243 }
244
245 const ERROR_ACCESS_DENIED: i32 = 5;
255 const ERROR_SHARING_VIOLATION: i32 = 32;
256 const MAX_ATTEMPTS: u32 = 10;
257 let mut backoff = std::time::Duration::from_millis(1);
258 for attempt in 1..=MAX_ATTEMPTS {
259 if unsafe { MoveFileExW(temp_wide.as_ptr(), destination_wide.as_ptr(), flags) } != 0 {
262 return Ok(durability == AtomicWriteDurability::Flush);
263 }
264 let error = io::Error::last_os_error();
265 let retryable = matches!(
266 error.raw_os_error(),
267 Some(ERROR_SHARING_VIOLATION | ERROR_ACCESS_DENIED)
268 );
269 if !retryable || attempt == MAX_ATTEMPTS {
270 return Err(error);
271 }
272 std::thread::sleep(backoff);
273 backoff = (backoff * 2).min(std::time::Duration::from_millis(50));
274 }
275 unreachable!("the loop returns on the final attempt")
276}
277
278fn sync_parent_dir(path: &Path) -> bool {
279 if let Some(parent) = path.parent() {
280 if parent.as_os_str().is_empty() {
281 return false;
282 }
283 if let Ok(dir) = OpenOptions::new().read(true).open(parent) {
284 return dir.sync_all().is_ok();
285 }
286 }
287 false
288}
289
290#[cfg(unix)]
293fn apply_mode(path: &Path, mode: u32) -> io::Result<()> {
294 use std::os::unix::fs::PermissionsExt;
295 std::fs::set_permissions(path, std::fs::Permissions::from_mode(mode))
296}
297
298#[cfg(not(unix))]
299fn apply_mode(_path: &Path, _mode: u32) -> io::Result<()> {
300 Ok(())
301}
302
303struct TempFile {
306 path: PathBuf,
307 file: Option<File>,
308}
309
310const TEMP_STEM_MAX: usize = 16;
312
313fn temp_sibling_name(file_name: &str) -> String {
325 let stem: String = file_name.chars().take(TEMP_STEM_MAX).collect();
326 format!(".{stem}.{}.tmp", uuid::Uuid::now_v7().simple())
327}
328
329impl TempFile {
330 fn create(target: &Path, mode: Option<u32>) -> io::Result<Self> {
331 let parent = target.parent().ok_or_else(|| {
332 io::Error::new(
333 io::ErrorKind::InvalidInput,
334 format!(
335 "atomic_io: destination '{}' has no parent directory",
336 target.display()
337 ),
338 )
339 })?;
340 if !parent.as_os_str().is_empty() {
341 std::fs::create_dir_all(parent)?;
342 }
343 let file_name = target
344 .file_name()
345 .and_then(|value| value.to_str())
346 .unwrap_or("file");
347 let tmp_name = temp_sibling_name(file_name);
348 let tmp_path = if parent.as_os_str().is_empty() {
349 PathBuf::from(tmp_name)
350 } else {
351 parent.join(tmp_name)
352 };
353 let file = OpenOptions::new()
354 .create_new(true)
355 .write(true)
356 .open(&tmp_path)?;
357 if let Some(mode) = mode {
358 if let Err(error) = apply_mode(&tmp_path, mode) {
359 drop(file);
360 let _ = std::fs::remove_file(&tmp_path);
361 return Err(error);
362 }
363 } else if let Ok(metadata) = std::fs::metadata(target) {
364 if let Err(error) = file.set_permissions(metadata.permissions()) {
365 drop(file);
366 let _ = std::fs::remove_file(&tmp_path);
367 return Err(error);
368 }
369 }
370 Ok(Self {
371 path: tmp_path,
372 file: Some(file),
373 })
374 }
375}
376
377#[cfg(test)]
378mod tests {
379 use super::*;
380
381 #[test]
382 fn temp_sibling_name_is_length_bounded_regardless_of_target_name() {
383 let bound = 1 + TEMP_STEM_MAX + 1 + 32 + 4; for name in ["s", "state.json", &"a".repeat(64), &"z".repeat(4096)] {
388 let temp = temp_sibling_name(name);
389 assert!(
390 temp.len() <= bound,
391 "temp name {:?} (len {}) exceeds bound {bound}",
392 temp,
393 temp.len()
394 );
395 assert!(temp.starts_with('.') && temp.ends_with(".tmp"));
396 }
397 }
398
399 #[test]
400 fn atomic_write_succeeds_for_a_long_target_file_name() {
401 let dir = tempfile::tempdir().unwrap();
404 let path = dir.path().join("a".repeat(200));
405 atomic_write(&path, b"payload").unwrap();
406 assert_eq!(std::fs::read(&path).unwrap(), b"payload");
407 }
408
409 #[test]
410 fn writes_bytes_atomically() {
411 let dir = tempfile::tempdir().unwrap();
412 let path = dir.path().join("state.json");
413 atomic_write(&path, b"hello").unwrap();
414 assert_eq!(std::fs::read(&path).unwrap(), b"hello");
415 }
416
417 #[test]
418 fn overwrites_existing_file() {
419 let dir = tempfile::tempdir().unwrap();
420 let path = dir.path().join("state.json");
421 std::fs::write(&path, b"old").unwrap();
422 atomic_write(&path, b"new").unwrap();
423 assert_eq!(std::fs::read(&path).unwrap(), b"new");
424 }
425
426 #[test]
427 fn creates_missing_parent_dirs() {
428 let dir = tempfile::tempdir().unwrap();
429 let path = dir.path().join("a/b/c/state.json");
430 atomic_write(&path, b"deep").unwrap();
431 assert_eq!(std::fs::read(&path).unwrap(), b"deep");
432 }
433
434 #[test]
435 fn streaming_writer_finalizes_atomically() {
436 let dir = tempfile::tempdir().unwrap();
437 let path = dir.path().join("log.jsonl");
438 atomic_write_with(&path, |writer| {
439 writeln!(writer, "first")?;
440 writeln!(writer, "second")?;
441 Ok(())
442 })
443 .unwrap();
444 let read = std::fs::read_to_string(&path).unwrap();
445 assert_eq!(read, "first\nsecond\n");
446 }
447
448 #[test]
449 fn streaming_writer_cleans_up_on_error() {
450 let dir = tempfile::tempdir().unwrap();
451 let path = dir.path().join("state.json");
452 std::fs::write(&path, b"old").unwrap();
453 let err = atomic_write_with(&path, |writer| {
454 writer.write_all(b"partial")?;
455 Err(io::Error::other("nope"))
456 })
457 .unwrap_err();
458 assert_eq!(err.to_string(), "nope");
459 assert_eq!(std::fs::read(&path).unwrap(), b"old");
460 let leftover: Vec<_> = std::fs::read_dir(dir.path())
462 .unwrap()
463 .filter_map(Result::ok)
464 .filter(|entry| entry.file_name().to_string_lossy().ends_with(".tmp"))
465 .collect();
466 assert!(
467 leftover.is_empty(),
468 "tmp file should be cleaned up on error"
469 );
470 }
471
472 #[test]
473 fn flush_and_replace_failures_preserve_destination_and_clean_up() {
474 for stage in ["flush", "replace"] {
475 let dir = tempfile::tempdir().unwrap();
476 let path = dir.path().join("state.json");
477 std::fs::write(&path, b"old").unwrap();
478 TEST_FAILURE_STAGE.with(|value| value.set(Some(stage)));
479 let error = atomic_write(&path, b"new").unwrap_err();
480 TEST_FAILURE_STAGE.with(|value| value.set(None));
481
482 assert_eq!(error.to_string(), format!("injected {stage} failure"));
483 assert_eq!(std::fs::read(&path).unwrap(), b"old");
484 let leftovers: Vec<_> = std::fs::read_dir(dir.path())
485 .unwrap()
486 .filter_map(Result::ok)
487 .filter(|entry| entry.file_name().to_string_lossy().ends_with(".tmp"))
488 .collect();
489 assert!(leftovers.is_empty(), "{stage} left a temp file");
490 }
491 }
492
493 #[cfg(unix)]
494 #[test]
495 fn replacement_preserves_existing_permissions() {
496 use std::os::unix::fs::PermissionsExt;
497
498 let dir = tempfile::tempdir().unwrap();
499 let path = dir.path().join("state.json");
500 std::fs::write(&path, b"old").unwrap();
501 std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o640)).unwrap();
502 atomic_write(&path, b"new").unwrap();
503 assert_eq!(
504 std::fs::metadata(path).unwrap().permissions().mode() & 0o777,
505 0o640
506 );
507 }
508
509 #[cfg(unix)]
510 #[test]
511 fn mode_is_applied_before_the_rename() {
512 use std::os::unix::fs::PermissionsExt;
513 let dir = tempfile::tempdir().unwrap();
514 let path = dir.path().join("credentials.json");
515 atomic_write_with_mode(&path, b"secret", 0o600).unwrap();
516 let mode = std::fs::metadata(&path).unwrap().permissions().mode();
517 assert_eq!(mode & 0o777, 0o600, "credentials must be owner-only");
518 }
519
520 #[cfg(unix)]
521 #[test]
522 fn mode_survives_overwriting_a_loose_destination() {
523 use std::os::unix::fs::PermissionsExt;
524 let dir = tempfile::tempdir().unwrap();
525 let path = dir.path().join("credentials.json");
526 std::fs::write(&path, b"old").unwrap();
527 std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o644)).unwrap();
528 atomic_write_with_mode(&path, b"secret", 0o600).unwrap();
529 let mode = std::fs::metadata(&path).unwrap().permissions().mode();
530 assert_eq!(mode & 0o777, 0o600);
531 }
532
533 #[test]
534 fn concurrent_writers_do_not_collide() {
535 let dir = tempfile::tempdir().unwrap();
536 let path = std::sync::Arc::new(dir.path().join("state.json"));
537 let mut handles = Vec::new();
538 for i in 0..16 {
539 let path = std::sync::Arc::clone(&path);
540 handles.push(std::thread::spawn(move || {
541 let payload = format!("writer-{i}");
542 atomic_write(&path, payload.as_bytes()).unwrap();
543 }));
544 }
545 for handle in handles {
546 handle.join().unwrap();
547 }
548 let final_contents = std::fs::read_to_string(&*path).unwrap();
551 assert!(
552 final_contents.starts_with("writer-") && final_contents.len() <= "writer-15".len(),
553 "unexpected final contents: {final_contents:?}"
554 );
555 }
556}