From 1791e3fe3c288be00ea47de9829bee8e5bc2cf97 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?F=C3=A9lix=20Boisselier?= <felix@fboisselier.fr>
Date: Sun, 15 Sep 2024 17:12:08 +0200
Subject: [PATCH] offload the finish measurement and get sample in a separate
 thread

---
 shaketune/helpers/accelerometer.py | 47 +++++++++++++++++++++++++-----
 1 file changed, 40 insertions(+), 7 deletions(-)

diff --git a/shaketune/helpers/accelerometer.py b/shaketune/helpers/accelerometer.py
index 6ebb07e..5f0b6ca 100644
--- a/shaketune/helpers/accelerometer.py
+++ b/shaketune/helpers/accelerometer.py
@@ -16,6 +16,7 @@
 import time
 from multiprocessing import Process
 from pathlib import Path
+from threading import Thread
 from typing import List, Tuple, TypedDict
 
 import numpy as np
@@ -140,7 +141,7 @@ def wait_for_file_writes(self, k_reactor, timeout: int = 20):
 
         if not complete:
             raise TimeoutError(
-                'Shake&Tune was unable to write the accelerometer data into the archive file. '
+                'Shake&Tune was unable to write the accelerometer data into the .STDATA file. '
                 'This might be due to a slow SD card or a busy or full filesystem.'
             )
 
@@ -152,6 +153,7 @@ def __init__(self, klipper_accelerometer):
         self._k_accelerometer = klipper_accelerometer
         self._bg_client = None
         self._measurements_manager: MeasurementsManager = None
+        self._get_samples_thread = None
 
     @staticmethod
     def find_axis_accelerometer(printer, axis: str = 'xy'):
@@ -188,11 +190,42 @@ def stop_recording(self) -> MeasurementsManager:
 
         bg_client = self._bg_client
         self._bg_client = None
-        bg_client.finish_measurements()
 
-        samples = bg_client.samples or bg_client.get_samples()
-        self._measurements_manager.append_samples_to_last_measurement(samples)
-        m_manager = self._measurements_manager
-        self._measurements_manager = None
+        # Start a thread to finish the measurements and get the samples without blocking the main
+        # Klipper code. Since these operations are I/O-bound, the GIL should not cause significant
+        # issues and a thread is even recommended over a Process since bg_client may not be pickleable.
+        self._get_samples_thread = Thread(target=self._finish_and_get_samples, args=(bg_client,))
+        self._get_samples_thread.daemon = True
+        self._get_samples_thread.start()
 
-        return m_manager
+        return self._measurements_manager
+
+    def _finish_and_get_samples(self, bg_client):
+        try:
+            bg_client.finish_measurements()
+            samples = bg_client.samples or bg_client.get_samples()
+            self._measurements_manager.append_samples_to_last_measurement(samples)
+        except Exception as e:
+            ConsoleOutput.print(f'Error during accelerometer data retrieval: {e}')
+
+    def wait_for_samples(self, k_reactor, timeout: int = 20):
+        if self._get_samples_thread is None:
+            return  # No samples are being retrieved
+
+        eventtime = k_reactor.monotonic()
+        endtime = eventtime + timeout
+        complete = False
+
+        while eventtime < endtime:
+            eventtime = k_reactor.pause(eventtime + 0.05)
+            if not self._get_samples_thread.is_alive():
+                complete = True
+                break
+
+        if not complete:
+            raise TimeoutError(
+                'Shake&Tune was unable to retrieve accelerometer data in time. '
+                'This might be due to slow hardware or a busy system.'
+            )
+
+        self._get_samples_thread = None