Protect against partial writes by using write + rename pattern

1 files changed, 13 insertions, 4 deletions

diff --git a/femtoqueue.py b/femtoqueue.py
index 500089c..1d08290 100644
--- a/femtoqueue.py
+++ b/femtoqueue.py
@@ -15,7 +15,8 @@ class FemtoQueue:
         node_name: str,
         timeout_stale_ms: int = 30_000,
     ):
-        assert node_name != "pending" \
+        assert node_name != "creating" \
+           and node_name != "pending" \
            and node_name != "done" \
            and node_name != "failed"
         self.node_name = node_name
@@ -25,25 +26,33 @@ class FemtoQueue:
         self.latest_stale_check_ts: float | None = None
 
         self.data_dir = data_dir
+        self.dir_creating = path.join(data_dir, "creating")
         self.dir_pending = path.join(data_dir, "pending")
         self.dir_in_progress = path.join(data_dir, node_name)
         self.dir_done = path.join(data_dir, "done")
         self.dir_failed = path.join(data_dir, "failed")
 
         makedirs(self.data_dir, exist_ok=True)
+        makedirs(self.dir_creating, exist_ok=True)
         makedirs(self.dir_pending, exist_ok=True)
         makedirs(self.dir_in_progress, exist_ok=True)
         makedirs(self.dir_done, exist_ok=True)
         makedirs(self.dir_failed, exist_ok=True)
 
-    def push(self, data: bytes) -> str:
+    def _gen_increasing_uuid(self) -> str:
         now_us = 1_000_000 * time.time()
-        id = f"{str(int(now_us))}_{uuid4().hex[-12:]}"
+        return f"{str(int(now_us))}_{uuid4().hex[-12:]}"
+
+    def push(self, data: bytes) -> str:
+        id = self._gen_increasing_uuid()
+        creating_path = path.join(self.dir_creating, id)
         pending_path = path.join(self.dir_pending, id)
 
-        with open(pending_path, "wb") as f:
+        with open(creating_path, "wb") as f:
             f.write(data)
 
+        rename(creating_path, pending_path)
+
         return id
 
     def _release_stale_tasks(self):