python-pq 0.9.0__tar.gz → 0.9.1__tar.gz

{python_pq-0.9.0 → python_pq-0.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: python-pq
-Version: 0.9.0
+Version: 0.9.1
 Summary: Postgres-backed job queue for Python
 Keywords: postgres,job-queue,task-queue,background-jobs,worker
 Author: ricwo
@@ -73,7 +73,7 @@ Requires PostgreSQL and Python 3.13+.
 
 ## Features
 
-- **Fork isolation** -- Each task runs in a forked child process. If it OOMs, segfaults, or crashes, the worker keeps running.
+- **Fork isolation** -- Each task runs in a forked child process. If it OOMs, segfaults, or crashes, the worker keeps running. Each child gets a fresh connection pool and opens its own database connections.
 - **No extra infrastructure** -- Uses your existing Postgres. No broker to deploy, monitor, or lose data.
 - **Transactional enqueueing** -- Enqueue tasks in the same database transaction as your writes. If the transaction rolls back, the task is never created.
 - **Periodic tasks** -- Schedule with intervals (`timedelta`) or cron expressions. Control overlap, pause/resume without deleting.
@@ -261,6 +261,12 @@ pq.run_worker(pre_execute=setup_tracing, post_execute=flush_tracing)
 
 Hooks run in the forked child, making them safe for fork-unsafe resources like OpenTelemetry.
 
+### Calling pq from a task
+
+A task can use your application's `PQ` instance, for example to enqueue a follow-up task. Each forked child gets a fresh connection pool, opens its own connections, and closes them when the task ends. A child killed by its timeout or by a signal exits without closing them.
+
+If your own code forks (with `os.fork()` or `multiprocessing`), fork outside an open `pq.session()`. A connection checked out at fork time stays shared between the two processes.
+
 ## Logging
 
 pq uses [loguru](https://github.com/Delgan/loguru) internally and does **not** configure the logger on import. This means it inherits whatever logging setup the host application has already established.

{python_pq-0.9.0 → python_pq-0.9.1}/README.md

@@ -37,7 +37,7 @@ Requires PostgreSQL and Python 3.13+.
 
 ## Features
 
-- **Fork isolation** -- Each task runs in a forked child process. If it OOMs, segfaults, or crashes, the worker keeps running.
+- **Fork isolation** -- Each task runs in a forked child process. If it OOMs, segfaults, or crashes, the worker keeps running. Each child gets a fresh connection pool and opens its own database connections.
 - **No extra infrastructure** -- Uses your existing Postgres. No broker to deploy, monitor, or lose data.
 - **Transactional enqueueing** -- Enqueue tasks in the same database transaction as your writes. If the transaction rolls back, the task is never created.
 - **Periodic tasks** -- Schedule with intervals (`timedelta`) or cron expressions. Control overlap, pause/resume without deleting.
@@ -225,6 +225,12 @@ pq.run_worker(pre_execute=setup_tracing, post_execute=flush_tracing)
 
 Hooks run in the forked child, making them safe for fork-unsafe resources like OpenTelemetry.
 
+### Calling pq from a task
+
+A task can use your application's `PQ` instance, for example to enqueue a follow-up task. Each forked child gets a fresh connection pool, opens its own connections, and closes them when the task ends. A child killed by its timeout or by a signal exits without closing them.
+
+If your own code forks (with `os.fork()` or `multiprocessing`), fork outside an open `pq.session()`. A connection checked out at fork time stays shared between the two processes.
+
 ## Logging
 
 pq uses [loguru](https://github.com/Delgan/loguru) internally and does **not** configure the logger on import. This means it inherits whatever logging setup the host application has already established.

{python_pq-0.9.0 → python_pq-0.9.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "python-pq"
-version = "0.9.0"
+version = "0.9.1"
 description = "Postgres-backed job queue for Python"
 readme = "README.md"
 authors = [

{python_pq-0.9.0 → python_pq-0.9.1}/src/pq/__init__.py

@@ -5,7 +5,7 @@ from pq.models import Periodic, Task, TaskStatus
 from pq.priority import Priority
 from pq.worker import PostExecuteHook, PreExecuteHook, TaskTimeoutError
 
-__version__ = "0.9.0"
+__version__ = "0.9.1"
 
 __all__ = [
     "PQ",

{python_pq-0.9.0 → python_pq-0.9.1}/src/pq/client.py

@@ -2,6 +2,8 @@
 
 import importlib.resources
 import math
+import os
+import weakref
 from collections.abc import Callable, Set
 from contextlib import contextmanager
 from datetime import UTC, datetime, timedelta
@@ -62,6 +64,61 @@ def _validate_max_runtime(max_runtime_seconds: float | None) -> None:
         )
 
 
+_live_instances: "weakref.WeakSet[PQ]" = weakref.WeakSet()
+"""Every ``PQ`` instance, so a forked child can reset their pools. Weak
+references: the set never keeps an instance alive. ``close()`` does not remove
+an instance: its engine reconnects if it is used again."""
+
+
+def _reset_pools_after_fork() -> None:
+    """Give each ``PQ`` instance an empty pool in a forked child.
+
+    A child inherits the parent's pooled connections, sockets and TLS state
+    included. Querying on one makes two processes speak on one Postgres
+    session: over TLS the parent's next record fails with ``bad record mac``,
+    over a plain socket the two processes read each other's results.
+
+    ``close=False`` drops the inherited pool without closing its connections.
+    Closing them would send a terminate message on the parent's sockets and
+    break the parent's sessions.
+
+    Connections checked out at fork time are not covered: if the child returns
+    one (for example by leaving a ``with pq.session()`` block it was forked
+    in), the commit and rollback still go out on the parent's socket. Do not
+    fork inside an open ``pq.session()``.
+    """
+    _dispose_pools(
+        close=False, failure="Failed to reset a PQ connection pool after fork"
+    )
+
+
+def _close_pools_before_exit() -> None:
+    """Close the connections a forked child opened, before it calls
+    ``os._exit()``.
+
+    ``os._exit()`` skips interpreter cleanup, so without this Postgres sees
+    each connection drop without a terminate message: it counts the session
+    as abandoned, and over TLS logs ``unexpected eof while reading``. The
+    connections inherited from the parent are not in these pools any more
+    (``_reset_pools_after_fork`` dropped them), so this closes only the
+    child's own. Connections still checked out are left alone.
+    """
+    _dispose_pools(close=True, failure="Failed to close a PQ connection pool")
+
+
+def _dispose_pools(*, close: bool, failure: str) -> None:
+    for instance in list(_live_instances):
+        try:
+            instance._engine.dispose(close=close)
+        except Exception:
+            # One failing instance must not stop the others.
+            logger.exception(failure)
+
+
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=_reset_pools_after_fork)
+
+
 class PQ:
     """Postgres-backed task queue client."""
 
@@ -79,6 +136,7 @@ class PQ:
             pool_recycle=1800,
         )
         self._session_factory = sessionmaker(bind=self._engine)
+        _live_instances.add(self)
 
     @contextmanager
     def session(self) -> Any:

{python_pq-0.9.0 → python_pq-0.9.1}/src/pq/worker.py

@@ -26,6 +26,7 @@ from croniter import croniter
 from loguru import logger
 from sqlalchemy import func, or_, select, update
 
+from pq.client import _close_pools_before_exit
 from pq.models import Periodic, Task, TaskStatus
 from pq.registry import resolve_function_path
 from pq.serialization import deserialize
@@ -169,14 +170,20 @@ class _ShutdownState:
     ``run_worker`` resets it on entry.
     """
 
-    requested: bool = False
+    signum: int | None = None  # signal that requested the shutdown
     deadline: float | None = None  # time.monotonic() deadline, set on signal
     drain_timeout: float = DEFAULT_DRAIN_TIMEOUT
+    logged: bool = False  # log_request() already ran for this shutdown
+
+    @property
+    def requested(self) -> bool:
+        return self.signum is not None
 
     def reset(self, drain_timeout: float) -> None:
-        self.requested = False
+        self.signum = None
         self.deadline = None
         self.drain_timeout = drain_timeout
+        self.logged = False
 
     def deadline_passed(self) -> bool:
         """True once shutdown was requested and the drain budget is spent."""
@@ -184,25 +191,40 @@ class _ShutdownState:
             return False
         return self.deadline is None or time.monotonic() >= self.deadline
 
+    def log_request(self) -> None:
+        """Log the shutdown request once, from the worker loop.
+
+        The signal handler can't log it itself (see
+        ``_handle_shutdown_signal``), so the loops call this when they see
+        ``requested``.
+        """
+        if not self.requested or self.logged:
+            return
+        self.logged = True
+        logger.info(
+            f"Received signal {self.signum}, draining"
+            f" (drain_timeout={self.drain_timeout}s)..."
+        )
+
 
 _shutdown = _ShutdownState()
 
 
 def _handle_shutdown_signal(signum: int, frame: FrameType | None) -> None:
-    """SIGTERM/SIGINT handler: request a drain, never raise.
+    """SIGTERM/SIGINT handler: request a drain, never raise, never log.
 
     Raising from here would land at arbitrary bytecode and could orphan a
-    task (e.g. between claim-commit and child registration). Repeat signals
-    during the drain are ignored so they can't abandon remaining children.
+    task (e.g. between claim-commit and child registration). Logging from
+    here can deadlock: the handler runs on the main thread wherever it was
+    interrupted, and loguru holds its handler locks for the whole of
+    ``os.fork()``. The worker loops log the request instead
+    (``_ShutdownState.log_request``). Repeat signals during the drain are
+    ignored so they can't abandon remaining children.
     """
     if _shutdown.requested:
         return
     _shutdown.deadline = time.monotonic() + _shutdown.drain_timeout
-    _shutdown.requested = True
-    logger.info(
-        f"Received signal {signum}, draining"
-        f" (drain_timeout={_shutdown.drain_timeout}s)..."
-    )
+    _shutdown.signum = signum
 
 
 def _install_shutdown_handlers() -> dict[int, _SignalHandler] | None:
@@ -309,17 +331,23 @@ def _run_in_child(
             pass  # Don't let hook errors mask task errors
 
     if error is None:
-        os._exit(EXIT_SUCCESS)
+        exit_code = EXIT_SUCCESS
     elif isinstance(error, asyncio.TimeoutError):
-        os._exit(EXIT_TIMEOUT)
+        exit_code = EXIT_TIMEOUT
     else:
+        exit_code = EXIT_FAILURE
         # Send error message to parent via pipe
         try:
             error_msg = f"{type(error).__name__}: {error}\n{error_tb}"
             os.write(error_write_fd, error_msg.encode("utf-8", errors="replace"))
         except Exception:
             pass  # Best effort
-        os._exit(EXIT_FAILURE)
+
+    # Hang up cleanly on any connection the task opened through pq. Not done
+    # in _child_timeout_handler: a signal handler can interrupt pool code
+    # that holds the pool's lock.
+    _close_pools_before_exit()
+    os._exit(exit_code)
 
 
 def _fork_child(
@@ -385,6 +413,7 @@ def _wait_for_child(child_pid: int, read_fd: int) -> _ChildResult:
         pid, raw_status, rusage = os.wait4(child_pid, os.WNOHANG)
         if pid != 0:
             break
+        _shutdown.log_request()
         if not killed_by_drain and _shutdown.deadline_passed():
             _kill_child_group(child_pid)
             killed_by_drain = True
@@ -715,6 +744,7 @@ def run_worker(
         except KeyboardInterrupt:
             # Fallback for non-main-thread workers (no handlers installed).
             pass
+        _shutdown.log_request()
         logger.info("Worker stopped.")
     finally:
         _restore_shutdown_handlers(previous_handlers)
@@ -1488,6 +1518,7 @@ def _run_concurrent(
         # drain below still runs, just without a deadline.
         pass
 
+    _shutdown.log_request()
     if children:
         logger.info(f"Shutting down, draining {len(children)} in-flight task(s)...")
         # _wait_for_child inside _reap_and_update enforces the drain