fix: Resolve GitHub issues #76, #75, and #74 (#80)

* fix: Resolve GitHub issues #76, #75, and #74

Issue #76: Remove trailing slash from root_path in web UI
- SAQ's app.js expects root_path without trailing slash
- Prevents double-slash URLs like /saq//

Issue #75: Fix Redis async client detection
- Check both redis.asyncio.Redis and redis.Redis instances
- Supports pre-instantiated async Redis clients

Issue #74: Asyncio-first signal handling
- Use loop.add_signal_handler() instead of signal.signal()
- Handle both SIGTERM and SIGINT in worker subprocess
- Use asyncio.Event() for non-blocking shutdown coordination
- Properly cleanup signal handlers and cancel pending tasks

Closes #76, closes #75, closes #74

* fix: Respect shutdown_grace_period_s during worker termination

The parent process now waits for the configured grace period before
force-killing worker subprocesses. This ensures jobs have sufficient
time to complete gracefully when shutdown_grace_period_s is configured.

Changes:
- Add get_max_shutdown_timeout() to calculate timeout from workers
- Use maximum grace period + 2s buffer for termination timeout
- Display configured timeout in shutdown messages
- Update both CLI and plugin shutdown handlers

Author: cofin

litestar_saq/cli.py

@@ -1,10 +1,13 @@
+import contextlib
 import signal
 import sys
 import time
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional
 
 if TYPE_CHECKING:
+    import asyncio
     import multiprocessing
+    from collections.abc import Collection
 
     from click import Group
     from litestar import Litestar
@@ -13,16 +16,46 @@
     from litestar_saq.base import Worker
     from litestar_saq.plugin import SAQPlugin
 
+# Default timeout for graceful shutdown when no grace period is configured
+DEFAULT_SHUTDOWN_TIMEOUT = 5.0
+# Extra buffer time to allow for signal propagation and cleanup
+SHUTDOWN_BUFFER = 2.0
+
+
+def get_max_shutdown_timeout(workers: "Collection[Worker]") -> float:
+    """Calculate the maximum shutdown timeout from worker configurations.
+
+    The timeout is the maximum of all workers' shutdown_grace_period_s plus
+    a buffer for signal propagation. Falls back to DEFAULT_SHUTDOWN_TIMEOUT
+    if no grace periods are configured.
+
+    Args:
+        workers: Collection of worker instances.
+
+    Returns:
+        Maximum shutdown timeout in seconds.
+    """
+    grace_periods: list[float] = []
+    for worker in workers:
+        grace_period = getattr(worker, "_shutdown_grace_period_s", None)
+        if grace_period is not None:
+            grace_periods.append(grace_period)
+    if grace_periods:
+        return max(grace_periods) + SHUTDOWN_BUFFER
+    return DEFAULT_SHUTDOWN_TIMEOUT
+
 
 def _terminate_worker_processes(
     processes: "list[multiprocessing.Process]",
-    timeout: float = 5.0,
+    timeout: float = DEFAULT_SHUTDOWN_TIMEOUT,
 ) -> None:
     """Gracefully terminate worker processes with timeout.
 
     Args:
         processes: List of worker processes to terminate
-        timeout: Maximum time to wait for graceful shutdown in seconds
+        timeout: Maximum time to wait for graceful shutdown in seconds.
+            Should be at least as long as the worker's shutdown_grace_period_s
+            plus buffer time for signal propagation.
     """
     from litestar.cli._utils import console  # pyright: ignore
 
@@ -48,6 +81,101 @@ def _terminate_worker_processes(
                 console.print(f"[red]Error killing worker process: {p.name}[/]")
 
 
+def _get_event_loop() -> "asyncio.AbstractEventLoop":
+    import asyncio
+
+    try:
+        return asyncio.get_event_loop()
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        return loop
+
+
+def _register_shutdown_handlers(
+    loop: "asyncio.AbstractEventLoop",
+    shutdown_event: "asyncio.Event",
+) -> Callable[[], None]:
+    signal_handlers_registered = False
+    fallback_handlers_registered = False
+    original_sigterm: Any = None
+    original_sigint: Any = None
+
+    def request_shutdown() -> None:
+        shutdown_event.set()
+
+    try:
+        loop.add_signal_handler(signal.SIGTERM, request_shutdown)
+        loop.add_signal_handler(signal.SIGINT, request_shutdown)
+        signal_handlers_registered = True
+    except (NotImplementedError, RuntimeError):
+        original_sigterm = signal.getsignal(signal.SIGTERM)
+        original_sigint = signal.getsignal(signal.SIGINT)
+
+        def fallback_handler(_signum: int, _frame: Any) -> None:
+            loop.call_soon_threadsafe(shutdown_event.set)
+
+        try:
+            signal.signal(signal.SIGTERM, fallback_handler)
+            signal.signal(signal.SIGINT, fallback_handler)
+            fallback_handlers_registered = True
+        except ValueError:
+            fallback_handlers_registered = False
+
+    def cleanup() -> None:
+        if signal_handlers_registered:
+            loop.remove_signal_handler(signal.SIGTERM)
+            loop.remove_signal_handler(signal.SIGINT)
+            return
+        if not fallback_handlers_registered:
+            return
+        if original_sigterm is not None:
+            with contextlib.suppress(ValueError):
+                signal.signal(signal.SIGTERM, original_sigterm)
+        if original_sigint is not None:
+            with contextlib.suppress(ValueError):
+                signal.signal(signal.SIGINT, original_sigint)
+
+    return cleanup
+
+
+async def _run_worker_with_shutdown(worker: "Worker") -> None:
+    import asyncio
+
+    loop = asyncio.get_running_loop()
+    shutdown_event = asyncio.Event()
+    cleanup_handlers = _register_shutdown_handlers(loop, shutdown_event)
+
+    try:
+        await worker.queue.connect()
+        worker_task = asyncio.create_task(worker.start())
+        shutdown_task = asyncio.create_task(shutdown_event.wait())
+
+        done, pending = await asyncio.wait(
+            [worker_task, shutdown_task],
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+
+        if worker_task in done:
+            worker_task.result()
+
+        if shutdown_event.is_set():
+            await worker.stop()
+            if not worker_task.done():
+                worker_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await worker_task
+            pending.discard(worker_task)
+
+        for task in pending:
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+    finally:
+        await worker.queue.disconnect()
+        cleanup_handlers()
+
+
 def build_cli_app() -> "Group":  # noqa: C901, PLR0915
     import asyncio
     import multiprocessing
@@ -106,12 +234,13 @@ def run_worker(  # pyright: ignore[reportUnusedFunction]
         show_saq_info(app, workers, plugin)
         managed_workers = list(plugin.get_workers().values())
         processes: list[multiprocessing.Process] = []
+        shutdown_timeout = get_max_shutdown_timeout(managed_workers)
 
         def handle_shutdown_signal(signum: int, _frame: Any) -> None:
             """Handle shutdown signals (SIGTERM/SIGINT) for graceful shutdown."""
             sig_name = "SIGTERM" if signum == signal.SIGTERM else "SIGINT"
-            console.print(f"[yellow]Received {sig_name}, stopping workers...[/]")
-            _terminate_worker_processes(processes)
+            console.print(f"[yellow]Received {sig_name}, stopping workers (timeout: {shutdown_timeout:.1f}s)...[/]")
+            _terminate_worker_processes(processes, timeout=shutdown_timeout)
             loop = asyncio.get_event_loop()
             for w in managed_workers:
                 loop.run_until_complete(w.stop())
@@ -228,15 +357,9 @@ def run_saq_worker(worker: "Worker", logging_config: "Optional[BaseLoggingConfig
         worker: The worker instance to run.
         logging_config: Optional logging configuration to apply.
     """
-    import asyncio
-
     from litestar.logging.config import StructLoggingConfig
 
-    try:
-        loop = asyncio.get_event_loop()
-    except RuntimeError:
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
+    loop = _get_event_loop()
 
     if logging_config is not None:
         logging_config.configure()
@@ -246,24 +369,10 @@ def run_saq_worker(worker: "Worker", logging_config: "Optional[BaseLoggingConfig
     if worker.separate_process:
         worker.configure_structlog_context()
         if isinstance(logging_config, StructLoggingConfig) and logging_config.standard_lib_logging_config is not None:
-            _ = logging_config.standard_lib_logging_config.configure()
-
-    def handle_sigterm(_signum: int, _frame: Any) -> None:
-        """Handle SIGTERM in worker process."""
-        loop.run_until_complete(loop.create_task(worker.stop()))
-        sys.exit(0)
-
-    signal.signal(signal.SIGTERM, handle_sigterm)
-
-    async def worker_start(w: "Worker") -> None:
-        try:
-            await w.queue.connect()
-            await w.start()
-        finally:
-            await w.queue.disconnect()
+            logging_config.standard_lib_logging_config.configure()
 
     try:
         if worker.separate_process:
-            loop.run_until_complete(loop.create_task(worker_start(worker)))
+            loop.run_until_complete(_run_worker_with_shutdown(worker))
     except KeyboardInterrupt:
         loop.run_until_complete(loop.create_task(worker.stop()))

litestar_saq/config.py

@@ -427,7 +427,9 @@ def broker_type(self) -> 'Literal["redis", "postgres", "http"]':
         if self._broker_type is None and self.broker_instance is not None:
             if self._is_instance_of(self.broker_instance, "psycopg_pool", "AsyncConnectionPool"):
                 self._broker_type = "postgres"
-            elif self._is_instance_of(self.broker_instance, "redis", "Redis"):
+            elif self._is_instance_of(self.broker_instance, "redis.asyncio", "Redis") or self._is_instance_of(
+                self.broker_instance, "redis", "Redis"
+            ):
                 self._broker_type = "redis"
             elif self._is_instance_of(self.broker_instance, "saq.queue.http", "HttpQueue"):
                 self._broker_type = "http"
@@ -464,7 +466,9 @@ def queue_class(self) -> "type[Queue]":
                 from saq.queue.postgres import PostgresQueue
 
                 self._queue_class = PostgresQueue
-            elif self._is_instance_of(self.broker_instance, "redis", "Redis"):
+            elif self._is_instance_of(self.broker_instance, "redis.asyncio", "Redis") or self._is_instance_of(
+                self.broker_instance, "redis", "Redis"
+            ):
                 from saq.queue.redis import RedisQueue
 
                 self._queue_class = RedisQueue

litestar_saq/controllers.py

@@ -15,21 +15,34 @@
 
 
 async def job_info(queue: "Queue", job_id: str) -> "Job":
+    """Return a job
+
+    Raises:
+        NotFoundException: If the queue or job is not found.
+    """
     job = await queue.job(job_id)
     if not job:
         msg = f"Could not find job ID {job_id}"
         raise NotFoundException(msg)
     return cast("Job", job)
 
 
+def get_queue_or_404(task_queues: "TaskQueues", queue_id: str) -> "Queue":
+    """Return a queue or raise NotFoundException."""
+    queue = task_queues.queues.get(queue_id)
+    if queue is None:
+        msg = f"Could not find the {queue_id} queue"
+        raise NotFoundException(msg)
+    return queue
+
+
 @lru_cache(typed=True)
-def build_controller(  # noqa: C901
+def build_controller(
     url_base: str = "/saq",
     controller_guards: "Optional[list[Guard]]" = None,  # pyright: ignore[reportUnknownParameterType]
     include_in_schema_: bool = False,
 ) -> "type[Controller]":
     from litestar import Controller, MediaType, Response, get, post
-    from litestar.exceptions import NotFoundException
     from litestar.status_codes import HTTP_202_ACCEPTED, HTTP_500_INTERNAL_SERVER_ERROR
 
     normalized_root = url_base.rstrip("/") or "/saq"
@@ -45,7 +58,7 @@ def build_controller(  # noqa: C901
         </head>
         <body>
             <div id="app"></div>
-            <script>const root_path = "{escaped_root}/";</script>
+            <script>const root_path = "{escaped_root}";</script>
             <script src="{escaped_root}/static/snabbdom.js.gz"></script>
             <script src="{escaped_root}/static/app.js"></script>
         </body>
@@ -99,10 +112,7 @@ async def queue_detail(self, task_queues: "TaskQueues", queue_id: str) -> "dict[
             Returns:
                 The queue information.
             """
-            queue = task_queues.get(queue_id)
-            if not queue:
-                msg = f"Could not find the {queue_id} queue"
-                raise NotFoundException(msg)
+            queue = get_queue_or_404(task_queues, queue_id)
             return {"queue": await queue.info(jobs=True)}
 
         @get(
@@ -130,10 +140,7 @@ async def job_detail(
             Returns:
                 The job information.
             """
-            queue = task_queues.get(queue_id)
-            if not queue:
-                msg = f"Could not find the {queue_id} queue"
-                raise NotFoundException(msg)
+            queue = get_queue_or_404(task_queues, queue_id)
             job = await job_info(queue, job_id)
             job_dict = job.to_dict()
             if "kwargs" in job_dict:
@@ -160,16 +167,10 @@ async def job_retry(self, task_queues: "TaskQueues", queue_id: str, job_id: str)
                 queue_id: The queue ID.
                 job_id: The job ID.
 
-            Raises:
-                NotFoundException: If the queue or job is not found.
-
             Returns:
                 The job information.
             """
-            queue = task_queues.get(queue_id)
-            if not queue:
-                msg = f"Could not find the {queue_id} queue"
-                raise NotFoundException(msg)
+            queue = get_queue_or_404(task_queues, queue_id)
             job = await job_info(queue, job_id)
             await job.retry("retried from ui")
             return {}
@@ -198,10 +199,7 @@ async def job_abort(self, task_queues: "TaskQueues", queue_id: str, job_id: str)
             Returns:
                 The job information.
             """
-            queue = task_queues.get(queue_id)
-            if not queue:
-                msg = f"Could not find the {queue_id} queue"
-                raise NotFoundException(msg)
+            queue = get_queue_or_404(task_queues, queue_id)
             job = await job_info(queue, job_id)
             await job.abort("aborted from ui")
             return {}
@@ -236,7 +234,6 @@ async def health(self, task_queues: "TaskQueues") -> Response[str]:
                 status_code=HTTP_500_INTERNAL_SERVER_ERROR,
             )
 
-        # static site
         @get(
             [
                 url_base,

litestar_saq/decorators.py

@@ -170,6 +170,16 @@ async def _heartbeat_loop(job: "Optional[Job]", interval: float) -> None:
         return
 
     job_id = getattr(job, "id", "unknown")
+    logger.info(
+        "Heartbeat monitoring started for job %s (interval: %.1fs)",
+        job_id,
+        interval,
+        extra={
+            "job_id": job_id,
+            "event": "heartbeat_started",
+            "interval": interval,
+        },
+    )
 
     try:
         while True: