From 495c9ded58119c02ccc084feb24fa72b790fd5dc Mon Sep 17 00:00:00 2001
From: eupn <36292692+eupn@users.noreply.github.com>
Date: Fri, 26 Oct 2018 21:50:26 +0300
Subject: [PATCH] Use temporary file to save bitmaps (#1840)

---
 store/src/leaf_set.rs   | 18 +++++++++---------
 store/src/lib.rs        | 36 ++++++++++++++++++++++++++++++++++++
 store/src/prune_list.rs | 11 ++++++-----
 3 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/store/src/leaf_set.rs b/store/src/leaf_set.rs
index 0f44b752..a0febaf9 100644
--- a/store/src/leaf_set.rs
+++ b/store/src/leaf_set.rs
@@ -15,8 +15,6 @@
 //! Compact (roaring) bitmap representing the set of leaf positions
 //! that exist and are not currently pruned in the MMR.
 
-use std::fs::File;
-use std::io::{self, BufWriter, Read, Write};
 use std::path::Path;
 
 use croaring::Bitmap;
@@ -25,6 +23,10 @@ use core::core::hash::Hashed;
 use core::core::pmmr;
 use core::core::BlockHeader;
 use prune_list::PruneList;
+use save_via_temp_file;
+
+use std::fs::File;
+use std::io::{self, BufWriter, Read, Write};
 
 /// Compact (roaring) bitmap representing the set of positions of
 /// leaves that are currently unpruned in the MMR.
@@ -168,14 +170,12 @@ impl LeafSet {
 		// First run the optimization step on the bitmap.
 		self.bitmap.run_optimize();
 
-		// TODO - consider writing this to disk in a tmp file and then renaming?
-
 		// Write the updated bitmap file to disk.
-		{
-			let mut file = BufWriter::new(File::create(self.path.clone())?);
-			file.write_all(&self.bitmap.serialize())?;
-			file.flush()?;
-		}
+		save_via_temp_file(&self.path, ".tmp", |w| {
+			let mut w = BufWriter::new(w);
+			w.write_all(&self.bitmap.serialize())?;
+			w.flush()
+		})?;
 
 		// Make sure our backup in memory is up to date.
 		self.bitmap_bak = self.bitmap.clone();
diff --git a/store/src/lib.rs b/store/src/lib.rs
index 24aaf800..d7d51edc 100644
--- a/store/src/lib.rs
+++ b/store/src/lib.rs
@@ -76,3 +76,39 @@ pub fn u64_to_key<'a>(prefix: u8, val: u64) -> Vec<u8> {
 	u64_vec.insert(0, prefix);
 	u64_vec
 }
+
+/// Creates temporary file with name created by adding `temp_suffix` to `path`.
+/// Applies writer function to it and renames temporary file into original specified by `path`.
+pub fn save_via_temp_file<F>(
+	path: &str,
+	temp_suffix: &str,
+	mut writer: F,
+) -> Result<(), std::io::Error>
+where
+	F: FnMut(Box<std::io::Write>) -> Result<(), std::io::Error>,
+{
+	assert_ne!(*temp_suffix, *"");
+
+	use std::fs::{remove_file, rename, File};
+	use std::path::Path;
+
+	// Write temporary file
+	let temp_name = format!("{}{}", &path, temp_suffix);
+	let temp_path = Path::new(&temp_name);
+	if temp_path.exists() {
+		remove_file(&temp_path)?;
+	}
+
+	let file = File::create(&temp_path)?;
+	writer(Box::new(file))?;
+
+	// Move temporary file into original
+	let original = Path::new(&path);
+	if original.exists() {
+		remove_file(&original)?;
+	}
+
+	rename(&temp_path, &original)?;
+
+	Ok(())
+}
diff --git a/store/src/prune_list.rs b/store/src/prune_list.rs
index b61a1459..68017a15 100644
--- a/store/src/prune_list.rs
+++ b/store/src/prune_list.rs
@@ -28,6 +28,7 @@ use std::path::Path;
 use croaring::Bitmap;
 
 use core::core::pmmr::{bintree_postorder_height, family, path};
+use save_via_temp_file;
 
 /// Maintains a list of previously pruned nodes in PMMR, compacting the list as
 /// parents get pruned and allowing checking whether a leaf is pruned. Given
@@ -115,13 +116,13 @@ impl PruneList {
 		// Run the optimization step on the bitmap.
 		self.bitmap.run_optimize();
 
-		// TODO - consider writing this to disk in a tmp file and then renaming?
-
 		// Write the updated bitmap file to disk.
 		if let Some(ref path) = self.path {
-			let mut file = BufWriter::new(File::create(path)?);
-			file.write_all(&self.bitmap.serialize())?;
-			file.flush()?;
+			save_via_temp_file(&path, ".tmp", |w| {
+				let mut w = BufWriter::new(w);
+				w.write_all(&self.bitmap.serialize())?;
+				w.flush()
+			})?;
 		}
 
 		// Rebuild our "shift caches" here as we are flushing changes to disk