Atomic write of compressed files (#85)

* writes compressed files in a tmp file first

* adds test for tmp wkw writing

* tests

* ci

Author: normanrz

.github/workflows/python-module.yml

@@ -171,3 +171,20 @@ jobs:
           user: __token__
           password: ${{ secrets.PYPI_PASSWORD }}
           skip-existing: true
+
+  complete:
+    needs:
+      - test_rust
+      - build_lin
+      - build_mac
+      - build_win
+    if: always()
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check failure
+        if: |
+          contains(needs.*.result, 'failure') ||
+          contains(needs.*.result, 'cancelled')
+        run: exit 1
+      - name: Success
+        run: echo Success!

python/tests/test_wkw.py

@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 import wkw
 import numpy as np
@@ -384,6 +385,35 @@ def test_dataset_with_invalid_jumptable():
     assert "Corrupt jump table" in str(excinfo)
 
 
+def test_readwrite_compressed_tmp():
+    with wkw.Dataset.create(
+        "tests/tmp",
+        wkw.Header(np.uint8, block_type=wkw.Header.BLOCK_TYPE_LZ4, file_len=1),
+    ) as dataset:
+        test_data = generate_test_data(dataset.header.voxel_type)
+
+        wkw_path = path.join("tests/tmp", "z0", "y0", "x0.wkw")
+
+        # write data, should run through cleanly
+        dataset.write(POSITION, test_data)
+        assert np.array_equiv(dataset.read(POSITION, SIZE), test_data)
+        assert path.exists(wkw_path)
+        assert not path.exists(wkw_path + "_tmp")
+
+        # simulate broken write
+        os.rename(wkw_path, wkw_path + "_tmp")
+        wkw_file_bytes = bytearray(Path(wkw_path + "_tmp").read_bytes())
+        zeros = b"\x00" * 8
+        wkw_file_bytes[16 : (16 + 8)] = zeros
+        Path(wkw_path + "_tmp").write_bytes(wkw_file_bytes)
+
+        # try another write, should clean up broken write
+        dataset.write(POSITION, test_data)
+        assert path.exists(wkw_path)
+        assert not path.exists(wkw_path + "_tmp")
+        assert np.array_equiv(dataset.read(POSITION, SIZE), test_data)
+
+
 def generate_test_data(dtype, size=SIZE, order="C"):
     return np.array(
         np.random.uniform(np.iinfo(dtype).min, np.iinfo(dtype).max, size).astype(dtype),

rust/src/dataset.rs

@@ -1,6 +1,6 @@
 use std::fs;
 use std::path::{Path, PathBuf};
-use {BlockType, Box3, File, Header, Mat, Result, Vec3};
+use {Box3, File, Header, Mat, Result, Vec3};
 
 #[derive(Debug, Clone)]
 pub struct Dataset {
@@ -140,7 +140,7 @@ impl Dataset {
         }
 
         let file_len_vx_log2 = self.header.file_len_vx_log2() as u32;
-        if self.header.block_type == BlockType::LZ4 || self.header.block_type == BlockType::LZ4HC {
+        if self.header.is_compressed() {
             let file_len_vec = Vec3::from(1 << file_len_vx_log2);
             let is_dst_aligned = dst_pos % file_len_vec == Vec3::from(0);
             let is_shape_aligned = mat.shape % file_len_vec == Vec3::from(0);
@@ -170,6 +170,11 @@ impl Dataset {
                     cur_path.push(format!("y{}", cur_y));
                     cur_path.push(format!("x{}.wkw", cur_x));
 
+                    // writing compressed file into temporary file first
+                    if self.header.is_compressed() {
+                        cur_path.set_extension("wkw_tmp");
+                    }
+
                     // bounding box
                     let cur_file_ids = Vec3 {
                         x: cur_x,
@@ -186,22 +191,38 @@ impl Dataset {
                     let cur_src_pos = cur_box.min() - dst_pos;
                     let cur_dst_pos = cur_box.min() - cur_file_box.min();
 
-                    let mut file = match File::open_or_create(&cur_path, &self.header) {
-                        Ok(file) => file,
-                        Err(err) => {
-                            return Err(format!(
-                                "Error while open file {:?} for writing: {}",
-                                &cur_path, err
-                            ));
+                    {
+                        let mut file = match File::open_or_create(&cur_path, &self.header) {
+                            Ok(file) => file,
+                            Err(err) => {
+                                return Err(format!(
+                                    "Error while open file {:?} for writing: {}",
+                                    &cur_path, err
+                                ));
+                            }
+                        };
+                        match file.write_mat(cur_dst_pos, mat, cur_src_pos) {
+                            Ok(_) => {}
+                            Err(err) => {
+                                return Err(format!(
+                                    "Error while writing to file {:?}: {}",
+                                    &cur_path, err
+                                ));
+                            }
                         }
-                    };
-                    match file.write_mat(cur_dst_pos, mat, cur_src_pos) {
-                        Ok(_) => {}
-                        Err(err) => {
-                            return Err(format!(
-                                "Error while writing to file {:?}: {}",
-                                &cur_path, err
-                            ));
+                    }
+                    // moving compressed file into final file
+                    if self.header.is_compressed() {
+                        let mut new_path = cur_path.clone();
+                        new_path.set_extension("wkw");
+                        match File::rename(&cur_path, &new_path) {
+                            Ok(_) => {}
+                            Err(err) => {
+                                return Err(format!(
+                                    "Error while renaming temporary file {:?} to {:?}: {}",
+                                    &cur_path, &new_path, err
+                                ));
+                            }
                         }
                     }
                 }

rust/src/file.rs

@@ -51,7 +51,7 @@ impl File {
 
         let mut file = open_opts
             .open(path)
-            .or(Err(format!("Could not open file {:?}", path)))?;
+            .or(Err(format!("Could not open file {:?}", &path)))?;
 
         // check if file was created
         let (header, created) = match Header::read(&mut file) {
@@ -70,6 +70,10 @@ impl File {
         Ok(file)
     }
 
+    pub(crate) fn rename(old_path: &path::Path, new_path: &path::Path) -> Result<()> {
+        fs::rename(old_path, new_path).or_else(|err| Err(err.to_string()))
+    }
+
     pub(crate) fn read_mat(
         &mut self,
         src_pos: Vec3,

rust/src/header.rs

@@ -250,6 +250,10 @@ impl Header {
         self.voxel_size as usize > self.voxel_type_size()
     }
 
+    pub fn is_compressed(&self) -> bool {
+        return self.block_type == BlockType::LZ4 || self.block_type == BlockType::LZ4HC;
+    }
+
     fn from_bytes(buf: [u8; 16]) -> Result<Header> {
         let raw: HeaderRaw = unsafe { mem::transmute(buf) };