🪲 BUG-#140: Fix race condition, thread tracking, and signal handler restoration in SchedulerCron

dotflow/providers/scheduler_cron.py

@@ -92,6 +92,9 @@ def __init__(
         self._lock = threading.Lock()
         self._queue_count = 0
         self._parallel_semaphore = threading.Semaphore(10)
+        self._threads: list[threading.Thread] = []
+        self._prev_sigint = None
+        self._prev_sigterm = None
 
     def start(self, workflow: Callable, **kwargs) -> None:
         """Start the scheduler loop. Blocks the main thread.
@@ -123,9 +126,17 @@ def start(self, workflow: Callable, **kwargs) -> None:
 
             self._dispatch(workflow=workflow, **kwargs)
 
-    def stop(self) -> None:
-        """Stop the scheduler loop gracefully."""
+    def stop(self, timeout: float | None = None) -> None:
+        """Stop the scheduler loop and wait for in-flight threads.
+
+        Args:
+            timeout: Max seconds to wait for each thread. None = wait forever.
+        """
         self.running = False
+        self._restore_signals()
+        for thread in self._threads:
+            thread.join(timeout=timeout)
+        self._threads.clear()
 
     def _dispatch(self, workflow: Callable, **kwargs) -> None:
         if self.overlap == TypeOverlap.SKIP:
@@ -151,6 +162,7 @@ def _dispatch_skip(self, workflow: Callable, **kwargs) -> None:
             args=(workflow,),
             kwargs=kwargs,
         )
+        self._threads.append(thread)
         thread.start()
 
     def _dispatch_queue(self, workflow: Callable, **kwargs) -> None:
@@ -166,6 +178,7 @@ def _dispatch_queue(self, workflow: Callable, **kwargs) -> None:
             args=(workflow,),
             kwargs=kwargs,
         )
+        self._threads.append(thread)
         thread.start()
 
     def _dispatch_parallel(self, workflow: Callable, **kwargs) -> None:
@@ -177,6 +190,7 @@ def _dispatch_parallel(self, workflow: Callable, **kwargs) -> None:
             args=(workflow,),
             kwargs=kwargs,
         )
+        self._threads.append(thread)
         thread.start()
 
     def _execute_parallel(self, workflow: Callable, **kwargs) -> None:
@@ -204,21 +218,32 @@ def _execute_queued(self, workflow: Callable, **kwargs) -> None:
         finally:
             with self._lock:
                 if self._queue_count > 0:
-                    self._queue_count -= 1
-                    thread = threading.Thread(
+                    self._queue_count = 0
+                    next_thread = threading.Thread(
                         target=self._execute_queued,
                         args=(workflow,),
                         kwargs=kwargs,
                     )
-                    thread.start()
                 else:
                     self._executing = False
+                    next_thread = None
+
+            if next_thread is not None:
+                next_thread.start()
 
     def _register_signals(self) -> None:
         if threading.current_thread() is not threading.main_thread():
             return
-        signal.signal(signal.SIGINT, self._handle_signal)
-        signal.signal(signal.SIGTERM, self._handle_signal)
+        self._prev_sigint = signal.signal(signal.SIGINT, self._handle_signal)
+        self._prev_sigterm = signal.signal(signal.SIGTERM, self._handle_signal)
+
+    def _restore_signals(self) -> None:
+        if threading.current_thread() is not threading.main_thread():
+            return
+        if self._prev_sigint is not None:
+            signal.signal(signal.SIGINT, self._prev_sigint)
+        if self._prev_sigterm is not None:
+            signal.signal(signal.SIGTERM, self._prev_sigterm)
 
     def _handle_signal(self, signum, frame) -> None:
         self.stop()