diff --git a/dotflow/providers/scheduler_cron.py b/dotflow/providers/scheduler_cron.py index ce62e7f8..bdf243f5 100644 --- a/dotflow/providers/scheduler_cron.py +++ b/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()