Bugfix/interrupts (#26)

* Fixed keyboardinterrupt handling.

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

.flake8

@@ -6,6 +6,8 @@ extend-ignore = E203
 docstring_style=sphinx
 
 ignore =
+  ; `noqa` comments overuse ))))
+  WPS402,
   ; Found `%` string formatting
   WPS323,
   ; Found block variable overlap.

taskiq/cli/args.py

@@ -27,6 +27,7 @@ class TaskiqArgs:
     log_collector_format: str
     max_threadpool_threads: int
     no_parse: bool
+    shutdown_timeout: float
 
     @classmethod
     def from_cli(cls, args: Optional[List[str]] = None) -> "TaskiqArgs":  # noqa: WPS213
@@ -106,6 +107,12 @@ def from_cli(cls, args: Optional[List[str]] = None) -> "TaskiqArgs":  # noqa: WP
             type=int,
             help="Maximum number of threads for executing sync functions.",
         )
+        parser.add_argument(
+            "--shutdown-timeout",
+            type=float,
+            default=5,
+            help="Maximum amount of time for graceful broker's shutdown is seconds.",
+        )
 
         if args is None:
             namespace = parser.parse_args(args)

taskiq/cli/worker.py

@@ -126,7 +126,7 @@ def import_tasks(modules: list[str], pattern: str, fs_discover: bool) -> None:
     import_from_modules(modules)
 
 
-async def shutdown_broker(broker: AsyncBroker) -> None:
+async def shutdown_broker(broker: AsyncBroker, timeout: float) -> None:
     """
     This function used to shutdown broker.
 
@@ -136,11 +136,15 @@ async def shutdown_broker(broker: AsyncBroker) -> None:
     We need to handle such situations.
 
     :param broker: current broker.
+    :param timeout: maximum amout of time to shutdown the broker.
     """
+    logger.warning("Shutting down the broker.")
     try:
-        ret_val = await broker.shutdown()  # type: ignore
+        ret_val = await asyncio.wait_for(broker.shutdown(), timeout)  # type: ignore
         if ret_val is not None:
             logger.info("Broker returned value on shutdown: '%s'", str(ret_val))
+    except asyncio.TimeoutError:
+        logger.warning("Cannot shutdown broker gracefully. Timed out.")
     except Exception as exc:
         logger.warning(
             "Exception found while terminating: %s",
@@ -149,7 +153,7 @@ async def shutdown_broker(broker: AsyncBroker) -> None:
         )
 
 
-def start_listen(args: TaskiqArgs) -> None:
+def start_listen(args: TaskiqArgs) -> None:  # noqa: C901, WPS213
     """
     This function starts actual listening process.
 
@@ -172,14 +176,42 @@ def start_listen(args: TaskiqArgs) -> None:
     AsyncBroker.is_worker_process = True
     broker = import_broker(args.broker)
     import_tasks(args.modules, args.tasks_pattern, args.fs_discover)
-    loop = asyncio.get_event_loop()
     if not isinstance(broker, AsyncBroker):
         raise ValueError("Unknown broker type. Please use AsyncBroker instance.")
+
+    # Here how we manage interruptions.
+    # We have to remember shutting_down state,
+    # because KeyboardInterrupt can be send multiple
+    # times. And it may interrupt the broker's shutdown process.
+    shutting_down = False
+
+    def interrupt_handler(_signum: int, _frame: Any) -> None:
+        """
+        Signal handler.
+
+        This handler checks if process is already
+        terminating and if it's true, it does nothing.
+
+        :param _signum: received signal number.
+        :param _frame: current execution frame.
+        :raises KeyboardInterrupt: if termiation hasn't begun.
+        """
+        nonlocal shutting_down  # noqa: WPS420
+        if shutting_down:
+            return
+        shutting_down = True  # noqa: WPS442
+        raise KeyboardInterrupt
+
+    signal.signal(signal.SIGINT, interrupt_handler)
+
+    loop = asyncio.get_event_loop()
     try:
         loop.run_until_complete(async_listen_messages(broker, args))
-    except (KeyboardInterrupt, Exception):
-        logger.warning("Terminating process!")
-    loop.run_until_complete(broker.shutdown())
+    except KeyboardInterrupt:
+        logger.warning("Worker process interrupted.")
+    except Exception as exc:
+        logger.error("Exception found: %s", exc, exc_info=True)
+    loop.run_until_complete(shutdown_broker(broker, args.shutdown_timeout))
 
 
 def watch_workers_restarts(args: TaskiqArgs) -> None: