using reactor time instead of join to wait for S&T process to finish

Frix-x · Jun 17, 2024 · 9b0ff2c · 9b0ff2c
1 parent e4a6255
commit 9b0ff2c
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 12 deletions.
diff --git a/shaketune/shaketune.py b/shaketune/shaketune.py
@@ -117,29 +117,54 @@ def __init__(self, config) -> None:
     def cmd_EXCITATE_AXIS_AT_FREQ(self, gcmd) -> None:
         ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
         static_freq_graph_creator = StaticGraphCreator(self._config)
-        st_process = ShakeTuneProcess(self._config, static_freq_graph_creator, self.timeout)
+        st_process = ShakeTuneProcess(
+            self._config,
+            self._printer.get_reactor(),
+            static_freq_graph_creator,
+            self.timeout,
+        )
         excitate_axis_at_freq(gcmd, self._pconfig, st_process)
 
     def cmd_AXES_MAP_CALIBRATION(self, gcmd) -> None:
         ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
         axes_map_graph_creator = AxesMapGraphCreator(self._config)
-        st_process = ShakeTuneProcess(self._config, axes_map_graph_creator, self.timeout)
+        st_process = ShakeTuneProcess(
+            self._config,
+            self._printer.get_reactor(),
+            axes_map_graph_creator,
+            self.timeout,
+        )
         axes_map_calibration(gcmd, self._pconfig, st_process)
 
     def cmd_COMPARE_BELTS_RESPONSES(self, gcmd) -> None:
         ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
         belt_graph_creator = BeltsGraphCreator(self._config)
-        st_process = ShakeTuneProcess(self._config, belt_graph_creator, self.timeout)
+        st_process = ShakeTuneProcess(
+            self._config,
+            self._printer.get_reactor(),
+            belt_graph_creator,
+            self.timeout,
+        )
         compare_belts_responses(gcmd, self._pconfig, st_process)
 
     def cmd_AXES_SHAPER_CALIBRATION(self, gcmd) -> None:
         ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
         shaper_graph_creator = ShaperGraphCreator(self._config)
-        st_process = ShakeTuneProcess(self._config, shaper_graph_creator, self.timeout)
+        st_process = ShakeTuneProcess(
+            self._config,
+            self._printer.get_reactor(),
+            shaper_graph_creator,
+            self.timeout,
+        )
         axes_shaper_calibration(gcmd, self._pconfig, st_process)
 
     def cmd_CREATE_VIBRATIONS_PROFILE(self, gcmd) -> None:
         ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
         vibration_profile_creator = VibrationsGraphCreator(self._config)
-        st_process = ShakeTuneProcess(self._config, vibration_profile_creator, self.timeout)
+        st_process = ShakeTuneProcess(
+            self._config,
+            self._printer.get_reactor(),
+            vibration_profile_creator,
+            self.timeout,
+        )
         create_vibrations_profile(gcmd, self._pconfig, st_process)
diff --git a/shaketune/shaketune_process.py b/shaketune/shaketune_process.py
@@ -19,11 +19,11 @@
 
 
 class ShakeTuneProcess:
-    def __init__(self, config: ShakeTuneConfig, graph_creator, timeout: Optional[float] = None) -> None:
-        self._config = config
+    def __init__(self, st_config: ShakeTuneConfig, reactor, graph_creator, timeout: Optional[float] = None) -> None:
+        self._config = st_config
+        self._reactor = reactor
         self.graph_creator = graph_creator
         self._timeout = timeout
-
         self._process = None
 
     def get_graph_creator(self):
@@ -35,16 +35,25 @@ def run(self) -> None:
         self._process.start()
 
     def wait_for_completion(self) -> None:
-        if self._process is not None:
-            self._process.join()
+        if self._process is None:
+            return  # Nothing to wait for
+        eventtime = self._reactor.monotonic()
+        endtime = eventtime + self._timeout
+        complete = False
+        while eventtime < endtime:
+            eventtime = self._reactor.pause(eventtime + 0.05)
+            if not self._process.is_alive():
+                complete = True
+                break
+        if not complete:
+            self._handle_timeout()
 
     # This function is a simple wrapper to start the Shake&Tune process. It's needed in order to get the timeout
     # as a Timer in a thread INSIDE the Shake&Tune child process to not interfere with the main Klipper process
     def _shaketune_process_wrapper(self, graph_creator, timeout) -> None:
         if timeout is not None:
-            timer = threading.Timer(timeout, self._handle_timeout)
+            timer = threading.Timer(timeout + 5, self._handle_timeout)  # Add 5 seconds to the timeout for safety
             timer.start()
-
         try:
             self._shaketune_process(graph_creator)
         finally: