From c1b365fdf7f56cf05d84c7b095bebc12101a1c12 Mon Sep 17 00:00:00 2001
From: Anastasia Lubennikova <anastasia@neon.tech>
Date: Wed, 18 May 2022 14:29:01 +0300
Subject: [PATCH] Use temp filename while writing ImageLayer file

---
 .../src/layered_repository/delta_layer.rs     | 24 ++++++++---
 .../src/layered_repository/image_layer.rs     | 42 +++++++++++++++----
 2 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/pageserver/src/layered_repository/delta_layer.rs b/pageserver/src/layered_repository/delta_layer.rs
index 1c48f3def5..855e2a9172 100644
--- a/pageserver/src/layered_repository/delta_layer.rs
+++ b/pageserver/src/layered_repository/delta_layer.rs
@@ -420,6 +420,21 @@ impl DeltaLayer {
         }
     }
 
+    fn temp_path_for(
+        conf: &PageServerConf,
+        timelineid: ZTimelineId,
+        tenantid: ZTenantId,
+        key_start: Key,
+        lsn_range: Range<Lsn>,
+    ) -> PathBuf {
+        conf.timeline_path(&timelineid, &tenantid).join(format!(
+            "{}-XXX__{:016X}-{:016X}.temp",
+            key_start,
+            u64::from(lsn_range.start),
+            u64::from(lsn_range.end)
+        ))
+    }
+
     ///
     /// Open the underlying file and read the metadata into memory, if it's
     /// not loaded already.
@@ -607,12 +622,9 @@ impl DeltaLayerWriter {
         //
         // Note: This overwrites any existing file. There shouldn't be any.
         // FIXME: throw an error instead?
-        let path = conf.timeline_path(&timelineid, &tenantid).join(format!(
-            "{}-XXX__{:016X}-{:016X}.temp",
-            key_start,
-            u64::from(lsn_range.start),
-            u64::from(lsn_range.end)
-        ));
+        let path =
+            DeltaLayer::temp_path_for(conf, timelineid, tenantid, key_start, lsn_range.clone());
+
         let mut file = VirtualFile::create(&path)?;
         // make room for the header block
         file.seek(SeekFrom::Start(PAGE_SZ as u64))?;
diff --git a/pageserver/src/layered_repository/image_layer.rs b/pageserver/src/layered_repository/image_layer.rs
index c0c8e7789a..0a7cd2cdba 100644
--- a/pageserver/src/layered_repository/image_layer.rs
+++ b/pageserver/src/layered_repository/image_layer.rs
@@ -241,6 +241,20 @@ impl ImageLayer {
         }
     }
 
+    fn temp_path_for(
+        path_or_conf: &PathOrConf,
+        timelineid: ZTimelineId,
+        tenantid: ZTenantId,
+        fname: &ImageFileName,
+    ) -> PathBuf {
+        match path_or_conf {
+            PathOrConf::Path(path) => path.to_path_buf(),
+            PathOrConf::Conf(conf) => conf
+                .timeline_path(&timelineid, &tenantid)
+                .join(format!("{}.temp", fname)),
+        }
+    }
+
     ///
     /// Open the underlying file and read the metadata into memory, if it's
     /// not loaded already.
@@ -398,7 +412,7 @@ impl ImageLayer {
 ///
 pub struct ImageLayerWriter {
     conf: &'static PageServerConf,
-    _path: PathBuf,
+    path: PathBuf,
     timelineid: ZTimelineId,
     tenantid: ZTenantId,
     key_range: Range<Key>,
@@ -416,11 +430,9 @@ impl ImageLayerWriter {
         key_range: &Range<Key>,
         lsn: Lsn,
     ) -> anyhow::Result<ImageLayerWriter> {
-        // Create the file
-        //
-        // Note: This overwrites any existing file. There shouldn't be any.
-        // FIXME: throw an error instead?
-        let path = ImageLayer::path_for(
+        // Create the file initially with a temporary filename.
+        // We'll atomically rename it to the final name when we're done.
+        let path = ImageLayer::temp_path_for(
             &PathOrConf::Conf(conf),
             timelineid,
             tenantid,
@@ -441,7 +453,7 @@ impl ImageLayerWriter {
 
         let writer = ImageLayerWriter {
             conf,
-            _path: path,
+            path,
             timelineid,
             tenantid,
             key_range: key_range.clone(),
@@ -512,6 +524,22 @@ impl ImageLayerWriter {
                 index_root_blk,
             }),
         };
+
+        // Rename the file to its final name
+        //
+        // Note: This overwrites any existing file. There shouldn't be any.
+        // FIXME: throw an error instead?
+        let final_path = ImageLayer::path_for(
+            &PathOrConf::Conf(self.conf),
+            self.timelineid,
+            self.tenantid,
+            &ImageFileName {
+                key_range: self.key_range.clone(),
+                lsn: self.lsn,
+            },
+        );
+        std::fs::rename(self.path, &final_path)?;
+
         trace!("created image layer {}", layer.path().display());
 
         Ok(layer)