From 814321f524634821ad9868e4775a460345da175c Mon Sep 17 00:00:00 2001
From: Christos Falas <raspberrymaker@gmail.com>
Date: Mon, 27 Jan 2025 17:28:47 +0000
Subject: [PATCH] fix: Write to file race condition

---
 where_fi/cli/file.py | 38 +++++++++++++++++++++++++++++++-------
 1 file changed, 31 insertions(+), 7 deletions(-)

diff --git a/where_fi/cli/file.py b/where_fi/cli/file.py
index 4bcea7c..f195ce0 100644
--- a/where_fi/cli/file.py
+++ b/where_fi/cli/file.py
@@ -1,28 +1,52 @@
 import logging
+import multiprocessing as mp
+import threading
 from datetime import datetime
 from pathlib import Path
+from typing import Any
 
-import h5py
 import numpy as np
 import numpy.typing as npt
 import typer
 
 from . import globals
 
-
 app = typer.Typer()
 logger = logging.getLogger(__name__)
 
 
+def write_to_file(path: Path, queue: "mp.Queue[Any]") -> None:
+    import h5py
+
+    with h5py.File(path, "w") as file:
+        logger.info("Starting writer")
+        while True:
+            data = queue.get()
+            if data is None:
+                break
+            logger.info("Writing data to file")
+            file.create_dataset(datetime.now().isoformat(), data=data)
+        logger.info("Writer stopped")
+
+
 @app.command()
 def capture(output_path: Path) -> None:
     """Capture CSI data to a file"""
     logger.info(f"Capturing data to {output_path}")
 
-    with h5py.File(output_path, "w") as file:
+    queue: "mp.Queue[None | npt.NDArray[np.complex128]]" = mp.Queue()
+    writer = threading.Thread(
+        target=write_to_file,
+        args=(
+            output_path,
+            queue,
+        ),
+    )
+    writer.start()
 
-        def callback(antenna_data: npt.NDArray[np.complex128]) -> None:
-            logger.info(f"Got final CSI data with shape {antenna_data.shape}")
-            file.create_dataset(datetime.now().isoformat(), data=antenna_data)
+    def callback(antenna_data: npt.NDArray[np.complex128]) -> None:
+        logger.info(f"Got final CSI data with shape {antenna_data.shape}")
+        queue.put(antenna_data)
 
-        globals.csi_producer(csi_callback=callback)
+    globals.csi_producer(csi_callback=callback)
+    queue.put(None)