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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (25) hide show
  1. {python_pq-0.8.1 → python_pq-0.9.1}/PKG-INFO +21 -2
  2. {python_pq-0.8.1 → python_pq-0.9.1}/README.md +20 -1
  3. {python_pq-0.8.1 → python_pq-0.9.1}/pyproject.toml +1 -1
  4. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/__init__.py +1 -1
  5. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/client.py +153 -16
  6. python_pq-0.9.1/src/pq/migrations/versions/20261007T120000Z_d4e5f6a7b8c9_add_task_requeue.py +43 -0
  7. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/models.py +31 -0
  8. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/worker.py +145 -76
  9. {python_pq-0.8.1 → python_pq-0.9.1}/LICENSE +0 -0
  10. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/config.py +0 -0
  11. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/logging.py +0 -0
  12. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/migrations/README +0 -0
  13. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/migrations/__init__.py +0 -0
  14. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/migrations/env.py +0 -0
  15. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/migrations/script.py.mako +0 -0
  16. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/migrations/versions/20260109T055839Z_476683af098d_initial_schema.py +0 -0
  17. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/migrations/versions/20260109T063747Z_2483bec70083_add_client_id.py +0 -0
  18. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/migrations/versions/20260205T120000Z_a1b2c3d4e5f6_add_max_concurrent.py +0 -0
  19. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/migrations/versions/20260205T180000Z_b7c8d9e0f1a2_add_periodic_key.py +0 -0
  20. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/migrations/versions/20260217T120000Z_c3d4e5f6a7b8_add_periodic_active.py +0 -0
  21. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/migrations/versions/20260527T170136Z_aee3e8e7e647_add_task_max_runtime_seconds.py +0 -0
  22. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/migrations/versions/__init__.py +0 -0
  23. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/priority.py +0 -0
  24. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/registry.py +0 -0
  25. {python_pq-0.8.1 → python_pq-0.9.1}/src/pq/serialization.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: python-pq
3
- Version: 0.8.1
3
+ Version: 0.9.1
4
4
  Summary: Postgres-backed job queue for Python
5
5
  Keywords: postgres,job-queue,task-queue,background-jobs,worker
6
6
  Author: ricwo
@@ -73,7 +73,7 @@ Requires PostgreSQL and Python 3.13+.
73
73
 
74
74
  ## Features
75
75
 
76
- - **Fork isolation** -- Each task runs in a forked child process. If it OOMs, segfaults, or crashes, the worker keeps running.
76
+ - **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.
77
77
  - **No extra infrastructure** -- Uses your existing Postgres. No broker to deploy, monitor, or lose data.
78
78
  - **Transactional enqueueing** -- Enqueue tasks in the same database transaction as your writes. If the transaction rolls back, the task is never created.
79
79
  - **Periodic tasks** -- Schedule with intervals (`timedelta`) or cron expressions. Control overlap, pause/resume without deleting.
@@ -143,6 +143,19 @@ pq.upsert(send_email, to="a@b.com", client_id="welcome-email")
143
143
  pq.upsert(send_email, to="new@b.com", client_id="welcome-email")
144
144
  ```
145
145
 
146
+ If the task is running, the new version is queued and runs after the current run ends, so one `client_id` does not run twice at the same time. This lets a task schedule its own continuation:
147
+
148
+ ```python
149
+ def backfill(offset: int) -> None:
150
+ more = process_rows(offset, limit=1000)
151
+ if more:
152
+ with PQ(DATABASE_URL) as client:
153
+ # Runs after this run ends, never next to it
154
+ client.upsert(backfill, offset=offset + 1000, client_id="backfill")
155
+ ```
156
+
157
+ If the worker running the task dies, the queued version runs after the stale-task reaper picks up the row (`stale_task_timeout`, default 1 hour). Keep `stale_task_timeout` longer than your longest run: if the reaper picks up a run that is still alive, the queued version can start next to it.
158
+
146
159
  ## Periodic Tasks
147
160
 
148
161
  ### Intervals
@@ -248,6 +261,12 @@ pq.run_worker(pre_execute=setup_tracing, post_execute=flush_tracing)
248
261
 
249
262
  Hooks run in the forked child, making them safe for fork-unsafe resources like OpenTelemetry.
250
263
 
264
+ ### Calling pq from a task
265
+
266
+ 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.
267
+
268
+ 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.
269
+
251
270
  ## Logging
252
271
 
253
272
  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.
@@ -37,7 +37,7 @@ Requires PostgreSQL and Python 3.13+.
37
37
 
38
38
  ## Features
39
39
 
40
- - **Fork isolation** -- Each task runs in a forked child process. If it OOMs, segfaults, or crashes, the worker keeps running.
40
+ - **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.
41
41
  - **No extra infrastructure** -- Uses your existing Postgres. No broker to deploy, monitor, or lose data.
42
42
  - **Transactional enqueueing** -- Enqueue tasks in the same database transaction as your writes. If the transaction rolls back, the task is never created.
43
43
  - **Periodic tasks** -- Schedule with intervals (`timedelta`) or cron expressions. Control overlap, pause/resume without deleting.
@@ -107,6 +107,19 @@ pq.upsert(send_email, to="a@b.com", client_id="welcome-email")
107
107
  pq.upsert(send_email, to="new@b.com", client_id="welcome-email")
108
108
  ```
109
109
 
110
+ If the task is running, the new version is queued and runs after the current run ends, so one `client_id` does not run twice at the same time. This lets a task schedule its own continuation:
111
+
112
+ ```python
113
+ def backfill(offset: int) -> None:
114
+ more = process_rows(offset, limit=1000)
115
+ if more:
116
+ with PQ(DATABASE_URL) as client:
117
+ # Runs after this run ends, never next to it
118
+ client.upsert(backfill, offset=offset + 1000, client_id="backfill")
119
+ ```
120
+
121
+ If the worker running the task dies, the queued version runs after the stale-task reaper picks up the row (`stale_task_timeout`, default 1 hour). Keep `stale_task_timeout` longer than your longest run: if the reaper picks up a run that is still alive, the queued version can start next to it.
122
+
110
123
  ## Periodic Tasks
111
124
 
112
125
  ### Intervals
@@ -212,6 +225,12 @@ pq.run_worker(pre_execute=setup_tracing, post_execute=flush_tracing)
212
225
 
213
226
  Hooks run in the forked child, making them safe for fork-unsafe resources like OpenTelemetry.
214
227
 
228
+ ### Calling pq from a task
229
+
230
+ 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.
231
+
232
+ 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.
233
+
215
234
  ## Logging
216
235
 
217
236
  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.
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "python-pq"
3
- version = "0.8.1"
3
+ version = "0.9.1"
4
4
  description = "Postgres-backed job queue for Python"
5
5
  readme = "README.md"
6
6
  authors = [
@@ -5,7 +5,7 @@ from pq.models import Periodic, Task, TaskStatus
5
5
  from pq.priority import Priority
6
6
  from pq.worker import PostExecuteHook, PreExecuteHook, TaskTimeoutError
7
7
 
8
- __version__ = "0.8.1"
8
+ __version__ = "0.9.1"
9
9
 
10
10
  __all__ = [
11
11
  "PQ",
@@ -2,6 +2,8 @@
2
2
 
3
3
  import importlib.resources
4
4
  import math
5
+ import os
6
+ import weakref
5
7
  from collections.abc import Callable, Set
6
8
  from contextlib import contextmanager
7
9
  from datetime import UTC, datetime, timedelta
@@ -13,9 +15,9 @@ from croniter.croniter import CroniterBadCronError
13
15
  from loguru import logger
14
16
  import sqlalchemy as sa
15
17
  from sqlalchemy import create_engine, delete, func, select, update
16
- from sqlalchemy.dialects.postgresql import insert
18
+ from sqlalchemy.dialects.postgresql import JSONB, insert
17
19
  from sqlalchemy.engine import Engine
18
- from sqlalchemy.orm import sessionmaker
20
+ from sqlalchemy.orm import InstrumentedAttribute, sessionmaker
19
21
 
20
22
  from pq.models import Base, Periodic, Task, TaskStatus
21
23
  from pq.priority import Priority
@@ -62,6 +64,61 @@ def _validate_max_runtime(max_runtime_seconds: float | None) -> None:
62
64
  )
63
65
 
64
66
 
67
+ _live_instances: "weakref.WeakSet[PQ]" = weakref.WeakSet()
68
+ """Every ``PQ`` instance, so a forked child can reset their pools. Weak
69
+ references: the set never keeps an instance alive. ``close()`` does not remove
70
+ an instance: its engine reconnects if it is used again."""
71
+
72
+
73
+ def _reset_pools_after_fork() -> None:
74
+ """Give each ``PQ`` instance an empty pool in a forked child.
75
+
76
+ A child inherits the parent's pooled connections, sockets and TLS state
77
+ included. Querying on one makes two processes speak on one Postgres
78
+ session: over TLS the parent's next record fails with ``bad record mac``,
79
+ over a plain socket the two processes read each other's results.
80
+
81
+ ``close=False`` drops the inherited pool without closing its connections.
82
+ Closing them would send a terminate message on the parent's sockets and
83
+ break the parent's sessions.
84
+
85
+ Connections checked out at fork time are not covered: if the child returns
86
+ one (for example by leaving a ``with pq.session()`` block it was forked
87
+ in), the commit and rollback still go out on the parent's socket. Do not
88
+ fork inside an open ``pq.session()``.
89
+ """
90
+ _dispose_pools(
91
+ close=False, failure="Failed to reset a PQ connection pool after fork"
92
+ )
93
+
94
+
95
+ def _close_pools_before_exit() -> None:
96
+ """Close the connections a forked child opened, before it calls
97
+ ``os._exit()``.
98
+
99
+ ``os._exit()`` skips interpreter cleanup, so without this Postgres sees
100
+ each connection drop without a terminate message: it counts the session
101
+ as abandoned, and over TLS logs ``unexpected eof while reading``. The
102
+ connections inherited from the parent are not in these pools any more
103
+ (``_reset_pools_after_fork`` dropped them), so this closes only the
104
+ child's own. Connections still checked out are left alone.
105
+ """
106
+ _dispose_pools(close=True, failure="Failed to close a PQ connection pool")
107
+
108
+
109
+ def _dispose_pools(*, close: bool, failure: str) -> None:
110
+ for instance in list(_live_instances):
111
+ try:
112
+ instance._engine.dispose(close=close)
113
+ except Exception:
114
+ # One failing instance must not stop the others.
115
+ logger.exception(failure)
116
+
117
+
118
+ if hasattr(os, "register_at_fork"):
119
+ os.register_at_fork(after_in_child=_reset_pools_after_fork)
120
+
121
+
65
122
  class PQ:
66
123
  """Postgres-backed task queue client."""
67
124
 
@@ -79,6 +136,7 @@ class PQ:
79
136
  pool_recycle=1800,
80
137
  )
81
138
  self._session_factory = sessionmaker(bind=self._engine)
139
+ _live_instances.add(self)
82
140
 
83
141
  @contextmanager
84
142
  def session(self) -> Any:
@@ -230,6 +288,16 @@ class PQ:
230
288
  Behaves like enqueue(), but on conflict for client_id, updates all fields.
231
289
  Status resets to PENDING, attempts to 0, and timestamps are cleared.
232
290
 
291
+ If the existing task is RUNNING, it is not changed. The new version
292
+ is stored on the row and runs after the current run ends (the worker
293
+ or the stale reaper re-queues the row with it). A second upsert while
294
+ the row is still RUNNING replaces the stored version, so only the
295
+ latest one runs. One ``client_id`` does not run twice at the same
296
+ time, as long as the stale reaper's timeout is longer than the
297
+ longest run. If the worker running the task died, the stored version
298
+ runs only after the stale reaper re-queues the row: with the reaper
299
+ disabled (``stale_task_timeout=None``) it never runs.
300
+
233
301
  Args:
234
302
  task: Callable function to execute.
235
303
  *args: Positional arguments to pass to the handler.
@@ -255,6 +323,41 @@ class PQ:
255
323
  if run_at is None:
256
324
  run_at = datetime.now(UTC)
257
325
 
326
+ # In ``ON CONFLICT DO UPDATE``, ``Task.<column>`` refers to the
327
+ # existing row. A RUNNING row keeps every column as it is (the run
328
+ # in progress keeps its payload, ``started_at``, ``attempts`` and
329
+ # ``max_runtime_seconds``); the new version is parked in
330
+ # ``requeue`` and applied when the run ends. Any other status gets
331
+ # the plain overwrite. All in one statement, so the decision is
332
+ # made under the row lock that ``ON CONFLICT`` takes.
333
+ running = Task.status == TaskStatus.RUNNING
334
+ # ``run_at`` is added in SQL as ``timestamptz``, so Postgres
335
+ # serializes it as ISO 8601 with an offset. A naive ``run_at`` is
336
+ # then read in this session's time zone, the same as on the direct
337
+ # path, and not later in the session of whoever applies the version.
338
+ requeue = sa.literal(
339
+ {
340
+ "name": name,
341
+ "payload": payload,
342
+ "priority": int(priority),
343
+ "max_runtime_seconds": max_runtime_seconds,
344
+ },
345
+ JSONB,
346
+ ).op("||", return_type=JSONB)(
347
+ sa.func.jsonb_build_object(
348
+ "run_at", sa.cast(sa.literal(run_at), sa.DateTime(timezone=True))
349
+ )
350
+ )
351
+
352
+ def unless_running[T](
353
+ value: T, column: InstrumentedAttribute[T]
354
+ ) -> sa.ColumnElement[T]:
355
+ # ``sa.literal(..., column.type)`` binds the value through the
356
+ # column's type. Required for the ``task_status`` enum, which
357
+ # stores member names (``'PENDING'``): a bare enum member in a
358
+ # ``CASE`` would bind ``'pending'`` and fail.
359
+ return sa.case((running, column), else_=sa.literal(value, column.type))
360
+
258
361
  stmt = (
259
362
  insert(Task)
260
363
  .values(
@@ -269,16 +372,19 @@ class PQ:
269
372
  .on_conflict_do_update(
270
373
  index_elements=["client_id"],
271
374
  set_={
272
- "name": name,
273
- "payload": payload,
274
- "priority": priority,
275
- "status": TaskStatus.PENDING,
276
- "run_at": run_at,
277
- "max_runtime_seconds": max_runtime_seconds,
278
- "attempts": 0,
279
- "started_at": None,
280
- "completed_at": None,
281
- "error": None,
375
+ "name": unless_running(name, Task.name),
376
+ "payload": unless_running(payload, Task.payload),
377
+ "priority": unless_running(int(priority), Task.priority),
378
+ "status": unless_running(TaskStatus.PENDING, Task.status),
379
+ "run_at": unless_running(run_at, Task.run_at),
380
+ "max_runtime_seconds": unless_running(
381
+ max_runtime_seconds, Task.max_runtime_seconds
382
+ ),
383
+ "attempts": unless_running(0, Task.attempts),
384
+ "started_at": unless_running(None, Task.started_at),
385
+ "completed_at": unless_running(None, Task.completed_at),
386
+ "error": unless_running(None, Task.error),
387
+ "requeue": sa.case((running, requeue), else_=sa.null()),
282
388
  },
283
389
  )
284
390
  .returning(Task.id)
@@ -582,13 +688,17 @@ class PQ:
582
688
  return result.rowcount
583
689
 
584
690
  def reap_stale_tasks(self, threshold: timedelta) -> int:
585
- """Mark stale RUNNING tasks as FAILED.
691
+ """Mark stale RUNNING tasks as FAILED, or re-queue them.
586
692
 
587
693
  When a worker dies mid-execution (e.g. pod restart, OOM on the worker
588
694
  process), in-flight tasks stay RUNNING forever because no parent
589
695
  process remains to update their status. This method detects those
590
696
  orphaned rows and transitions them to FAILED.
591
697
 
698
+ A stale row that holds a version parked by ``upsert()`` (see
699
+ ``Task.requeue``) is re-queued as PENDING with that version instead
700
+ of being failed, so the parked version is not lost.
701
+
592
702
  Tasks enqueued with a per-task ``max_runtime_seconds`` override get a
593
703
  proportionally larger reaper window: the effective threshold for
594
704
  a row is ``max(threshold, max_runtime_seconds * 2)``. This avoids
@@ -605,7 +715,7 @@ class PQ:
605
715
  ``timedelta(seconds=max_runtime_seconds * 2)``.
606
716
 
607
717
  Returns:
608
- Number of tasks reaped.
718
+ Number of tasks reaped (failed plus re-queued).
609
719
  """
610
720
  now = datetime.now(UTC)
611
721
  threshold_seconds = threshold.total_seconds()
@@ -649,6 +759,7 @@ class PQ:
649
759
  .where(
650
760
  Task.status == TaskStatus.RUNNING,
651
761
  Task.started_at < stale_cutoff,
762
+ Task.requeue.is_(None),
652
763
  )
653
764
  .values(
654
765
  status=TaskStatus.FAILED,
@@ -669,7 +780,31 @@ class PQ:
669
780
  f" started_at={started_at},"
670
781
  f" max_runtime_seconds={max_runtime_seconds})"
671
782
  )
672
- return len(reaped)
783
+
784
+ # Stale rows with a parked upsert: re-queue them with the
785
+ # parked version. ``FOR UPDATE`` serializes with a concurrent
786
+ # ``upsert()`` (which may replace ``requeue``) and with the
787
+ # worker's own end-of-run update.
788
+ requeued = list(
789
+ session.execute(
790
+ select(Task)
791
+ .where(
792
+ Task.status == TaskStatus.RUNNING,
793
+ Task.started_at < stale_cutoff,
794
+ Task.requeue.is_not(None),
795
+ )
796
+ .with_for_update()
797
+ ).scalars()
798
+ )
799
+ for task in requeued:
800
+ logger.warning(
801
+ f"Reaped stale task '{task.name}' (id={task.id},"
802
+ f" started_at={task.started_at},"
803
+ f" max_runtime_seconds={task.max_runtime_seconds});"
804
+ " re-queued by a parked upsert"
805
+ )
806
+ task.apply_requeue()
807
+ return len(reaped) + len(requeued)
673
808
 
674
809
  def run_worker(
675
810
  self,
@@ -687,7 +822,9 @@ class PQ:
687
822
  On SIGTERM or SIGINT the worker stops claiming new tasks, waits up
688
823
  to ``drain_timeout`` seconds for in-flight tasks to finish, then
689
824
  kills any still-running task and marks its row FAILED with an
690
- explicit shutdown error. Interrupted tasks are not re-queued.
825
+ explicit shutdown error. Interrupted tasks are not re-queued,
826
+ unless an ``upsert()`` stored a new version while they ran: then the
827
+ row is re-queued with that version.
691
828
 
692
829
  Args:
693
830
  concurrency: Maximum number of tasks to process simultaneously.
@@ -0,0 +1,43 @@
1
+ """add requeue to pq_tasks
2
+
3
+ Revision ID: d4e5f6a7b8c9
4
+ Revises: aee3e8e7e647
5
+ Create Date: 2026-10-07 12:00:00 Z
6
+
7
+ Holds the next version of a one-off task when ``upsert()`` hits a row
8
+ that is RUNNING. The current run keeps its row unchanged; whoever ends
9
+ the run (worker or stale reaper) re-queues the row from this column
10
+ instead of closing it. This makes sure one ``client_id`` never runs
11
+ twice at the same time (ricwo/pq#27).
12
+
13
+ Nullable without a default, so on PostgreSQL 11+ ``ADD COLUMN`` is a
14
+ catalog-only change (no table rewrite). Existing rows get NULL, which
15
+ means "no parked version" — the previous behaviour.
16
+
17
+ """
18
+
19
+ from typing import Sequence, Union
20
+
21
+ from alembic import op
22
+ import sqlalchemy as sa
23
+ from sqlalchemy.dialects import postgresql
24
+
25
+
26
+ # revision identifiers, used by Alembic.
27
+ revision: str = "d4e5f6a7b8c9"
28
+ down_revision: Union[str, Sequence[str], None] = "aee3e8e7e647"
29
+ branch_labels: Union[str, Sequence[str], None] = None
30
+ depends_on: Union[str, Sequence[str], None] = None
31
+
32
+
33
+ def upgrade() -> None:
34
+ """Add the nullable ``requeue`` column to ``pq_tasks``."""
35
+ op.add_column(
36
+ "pq_tasks",
37
+ sa.Column("requeue", postgresql.JSONB(), nullable=True),
38
+ )
39
+
40
+
41
+ def downgrade() -> None:
42
+ """Remove the ``requeue`` column from ``pq_tasks``."""
43
+ op.drop_column("pq_tasks", "requeue")
@@ -79,6 +79,37 @@ class Task(Base):
79
79
  # ``max_runtime_seconds * 2``). Useful for occasionally-long tasks
80
80
  # in a fleet whose default is sized for typical short work.
81
81
  max_runtime_seconds: Mapped[float | None] = mapped_column(Float, nullable=True)
82
+ # Next version of this task, parked by an ``upsert()`` that arrived while
83
+ # the row was RUNNING: name, payload, priority, run_at (ISO 8601) and
84
+ # max_runtime_seconds. Whoever ends the current run (worker or reaper)
85
+ # re-queues the row with it instead of closing the row. NULL when no
86
+ # upsert is parked (the common case). ``none_as_null`` makes the ORM
87
+ # write SQL NULL for ``None`` (plain ``JSONB`` writes JSON ``null``,
88
+ # which ``IS NULL`` does not match).
89
+ requeue: Mapped[dict[str, Any] | None] = mapped_column(
90
+ JSONB(none_as_null=True), nullable=True
91
+ )
92
+
93
+ def apply_requeue(self) -> None:
94
+ """Turn this RUNNING row into a fresh PENDING task from ``requeue``.
95
+
96
+ The caller must hold a row lock (``FOR UPDATE``) and guarantee
97
+ that ``requeue`` is not None.
98
+ """
99
+ spec = self.requeue
100
+ if spec is None:
101
+ raise ValueError(f"Task id={self.id} has no parked requeue version")
102
+ self.name = spec["name"]
103
+ self.payload = spec["payload"]
104
+ self.priority = spec["priority"]
105
+ self.run_at = datetime.fromisoformat(spec["run_at"])
106
+ self.max_runtime_seconds = spec["max_runtime_seconds"]
107
+ self.status = TaskStatus.PENDING
108
+ self.attempts = 0
109
+ self.started_at = None
110
+ self.completed_at = None
111
+ self.error = None
112
+ self.requeue = None
82
113
 
83
114
 
84
115
  class Periodic(Base):
@@ -26,6 +26,7 @@ from croniter import croniter
26
26
  from loguru import logger
27
27
  from sqlalchemy import func, or_, select, update
28
28
 
29
+ from pq.client import _close_pools_before_exit
29
30
  from pq.models import Periodic, Task, TaskStatus
30
31
  from pq.registry import resolve_function_path
31
32
  from pq.serialization import deserialize
@@ -123,6 +124,9 @@ class _ChildSlot:
123
124
  name: str
124
125
  start_time: float
125
126
  is_periodic: bool
127
+ # ``started_at`` written by the claim of a one-off run. Identifies this
128
+ # run on the row (see ``_finish_one_off``). None for periodic slots.
129
+ started_at: datetime | None = None
126
130
  periodic_max_concurrent: int | None = None
127
131
 
128
132
 
@@ -166,14 +170,20 @@ class _ShutdownState:
166
170
  ``run_worker`` resets it on entry.
167
171
  """
168
172
 
169
- requested: bool = False
173
+ signum: int | None = None # signal that requested the shutdown
170
174
  deadline: float | None = None # time.monotonic() deadline, set on signal
171
175
  drain_timeout: float = DEFAULT_DRAIN_TIMEOUT
176
+ logged: bool = False # log_request() already ran for this shutdown
177
+
178
+ @property
179
+ def requested(self) -> bool:
180
+ return self.signum is not None
172
181
 
173
182
  def reset(self, drain_timeout: float) -> None:
174
- self.requested = False
183
+ self.signum = None
175
184
  self.deadline = None
176
185
  self.drain_timeout = drain_timeout
186
+ self.logged = False
177
187
 
178
188
  def deadline_passed(self) -> bool:
179
189
  """True once shutdown was requested and the drain budget is spent."""
@@ -181,25 +191,40 @@ class _ShutdownState:
181
191
  return False
182
192
  return self.deadline is None or time.monotonic() >= self.deadline
183
193
 
194
+ def log_request(self) -> None:
195
+ """Log the shutdown request once, from the worker loop.
196
+
197
+ The signal handler can't log it itself (see
198
+ ``_handle_shutdown_signal``), so the loops call this when they see
199
+ ``requested``.
200
+ """
201
+ if not self.requested or self.logged:
202
+ return
203
+ self.logged = True
204
+ logger.info(
205
+ f"Received signal {self.signum}, draining"
206
+ f" (drain_timeout={self.drain_timeout}s)..."
207
+ )
208
+
184
209
 
185
210
  _shutdown = _ShutdownState()
186
211
 
187
212
 
188
213
  def _handle_shutdown_signal(signum: int, frame: FrameType | None) -> None:
189
- """SIGTERM/SIGINT handler: request a drain, never raise.
214
+ """SIGTERM/SIGINT handler: request a drain, never raise, never log.
190
215
 
191
216
  Raising from here would land at arbitrary bytecode and could orphan a
192
- task (e.g. between claim-commit and child registration). Repeat signals
193
- during the drain are ignored so they can't abandon remaining children.
217
+ task (e.g. between claim-commit and child registration). Logging from
218
+ here can deadlock: the handler runs on the main thread wherever it was
219
+ interrupted, and loguru holds its handler locks for the whole of
220
+ ``os.fork()``. The worker loops log the request instead
221
+ (``_ShutdownState.log_request``). Repeat signals during the drain are
222
+ ignored so they can't abandon remaining children.
194
223
  """
195
224
  if _shutdown.requested:
196
225
  return
197
226
  _shutdown.deadline = time.monotonic() + _shutdown.drain_timeout
198
- _shutdown.requested = True
199
- logger.info(
200
- f"Received signal {signum}, draining"
201
- f" (drain_timeout={_shutdown.drain_timeout}s)..."
202
- )
227
+ _shutdown.signum = signum
203
228
 
204
229
 
205
230
  def _install_shutdown_handlers() -> dict[int, _SignalHandler] | None:
@@ -306,17 +331,23 @@ def _run_in_child(
306
331
  pass # Don't let hook errors mask task errors
307
332
 
308
333
  if error is None:
309
- os._exit(EXIT_SUCCESS)
334
+ exit_code = EXIT_SUCCESS
310
335
  elif isinstance(error, asyncio.TimeoutError):
311
- os._exit(EXIT_TIMEOUT)
336
+ exit_code = EXIT_TIMEOUT
312
337
  else:
338
+ exit_code = EXIT_FAILURE
313
339
  # Send error message to parent via pipe
314
340
  try:
315
341
  error_msg = f"{type(error).__name__}: {error}\n{error_tb}"
316
342
  os.write(error_write_fd, error_msg.encode("utf-8", errors="replace"))
317
343
  except Exception:
318
344
  pass # Best effort
319
- os._exit(EXIT_FAILURE)
345
+
346
+ # Hang up cleanly on any connection the task opened through pq. Not done
347
+ # in _child_timeout_handler: a signal handler can interrupt pool code
348
+ # that holds the pool's lock.
349
+ _close_pools_before_exit()
350
+ os._exit(exit_code)
320
351
 
321
352
 
322
353
  def _fork_child(
@@ -382,6 +413,7 @@ def _wait_for_child(child_pid: int, read_fd: int) -> _ChildResult:
382
413
  pid, raw_status, rusage = os.wait4(child_pid, os.WNOHANG)
383
414
  if pid != 0:
384
415
  break
416
+ _shutdown.log_request()
385
417
  if not killed_by_drain and _shutdown.deadline_passed():
386
418
  _kill_child_group(child_pid)
387
419
  killed_by_drain = True
@@ -569,7 +601,8 @@ def _maybe_reap_stale(
569
601
 
570
602
  Args:
571
603
  pq: PQ client instance.
572
- stale_task_timeout: RUNNING tasks older than this are marked FAILED.
604
+ stale_task_timeout: RUNNING tasks older than this are marked FAILED
605
+ (or re-queued, if an upsert parked a new version on them).
573
606
  ``None`` disables reaping.
574
607
  reaper_interval: Seconds between reaper checks.
575
608
  last_reap: Mutable list containing last reap timestamp.
@@ -581,9 +614,15 @@ def _maybe_reap_stale(
581
614
  if now - last_reap[0] < reaper_interval:
582
615
  return
583
616
 
584
- reaped = pq.reap_stale_tasks(stale_task_timeout)
585
- if reaped:
586
- logger.info(f"Reaped {reaped} stale RUNNING task(s)")
617
+ # A reaper error must not stop the worker loop (in concurrent mode that
618
+ # would also skip the drain of in-flight children). Retry next interval.
619
+ try:
620
+ reaped = pq.reap_stale_tasks(stale_task_timeout)
621
+ except Exception as e:
622
+ logger.error(f"Error reaping stale tasks: {e}")
623
+ else:
624
+ if reaped:
625
+ logger.info(f"Reaped {reaped} stale RUNNING task(s)")
587
626
 
588
627
  last_reap[0] = now
589
628
 
@@ -612,7 +651,9 @@ def run_worker(
612
651
  still-running task children and marks their rows FAILED with an explicit
613
652
  shutdown error. Interrupted tasks are NOT re-queued — pq's at-most-once
614
653
  semantics are preserved; applications that need redelivery must
615
- re-enqueue such tasks themselves.
654
+ re-enqueue such tasks themselves. The one exception: if an ``upsert()``
655
+ stored a new version on the row while it ran, the row is re-queued with
656
+ that new version (the interrupted run itself is not retried).
616
657
 
617
658
  Args:
618
659
  pq: PQ client instance.
@@ -625,9 +666,12 @@ def run_worker(
625
666
  retention_days: Days to keep completed/failed tasks. Default: 7.
626
667
  Set to 0 to disable automatic cleanup.
627
668
  cleanup_interval: Seconds between cleanup runs. Default: 3600 (1 hour).
628
- stale_task_timeout: Mark RUNNING tasks older than this as FAILED.
669
+ stale_task_timeout: Mark RUNNING tasks older than this as FAILED
670
+ (or re-queue them, if an upsert parked a new version on them).
629
671
  Catches orphaned tasks whose worker died mid-execution.
630
- Default: 1 hour. Set to ``None`` to disable.
672
+ Default: 1 hour. Set to ``None`` to disable. Must be longer
673
+ than the longest run: a live run that is reaped can overlap
674
+ with the version re-queued after it.
631
675
  drain_timeout: Seconds to wait for in-flight tasks when shutting
632
676
  down on SIGTERM/SIGINT. Default: 20. Must be set below the
633
677
  orchestrator's termination grace period, with headroom for the
@@ -700,6 +744,7 @@ def run_worker(
700
744
  except KeyboardInterrupt:
701
745
  # Fallback for non-main-thread workers (no handlers installed).
702
746
  pass
747
+ _shutdown.log_request()
703
748
  logger.info("Worker stopped.")
704
749
  finally:
705
750
  _restore_shutdown_handlers(previous_handlers)
@@ -756,6 +801,66 @@ def run_worker_once(
756
801
  return False
757
802
 
758
803
 
804
+ def _finish_one_off(
805
+ pq: PQ,
806
+ task_id: int,
807
+ started_at: datetime | None,
808
+ name: str,
809
+ status: TaskStatus,
810
+ error_msg: str | None,
811
+ ) -> None:
812
+ """Record the end of a one-off run.
813
+
814
+ Only touches the row if it is still RUNNING *with this run's*
815
+ ``started_at`` (the value written by its claim). So a reaper's FAILED
816
+ verdict (or a cancel) is not overwritten, and neither is a later run
817
+ of the same row: if the reaper re-queued the row while this run was
818
+ still alive and another worker claimed it, ``started_at`` no longer
819
+ matches. ``None`` never matches (a RUNNING row always has
820
+ ``started_at``), so the row is left alone. If an ``upsert()`` parked a
821
+ new version on the row while it ran (``Task.requeue``), the row is
822
+ re-queued as PENDING with that version instead of being closed. The
823
+ outcome of this run is then only in the worker log, the same as when
824
+ an upsert overwrites a COMPLETED or FAILED row.
825
+
826
+ ``FOR UPDATE`` serializes this read-then-write with a concurrent
827
+ ``upsert()``, whose ``ON CONFLICT DO UPDATE`` locks the same row: a
828
+ parked version is either seen here, or the upsert runs after this
829
+ commit and hits a non-RUNNING row, which it overwrites directly.
830
+ """
831
+ try:
832
+ with pq.session() as session:
833
+ task = session.execute(
834
+ select(Task)
835
+ .where(
836
+ Task.id == task_id,
837
+ Task.status == TaskStatus.RUNNING,
838
+ Task.started_at == started_at,
839
+ )
840
+ .with_for_update()
841
+ ).scalar_one_or_none()
842
+ if task is None:
843
+ logger.warning(
844
+ f"Task '{name}' (id={task_id}) was no longer RUNNING"
845
+ " for this run when the worker tried to record its result"
846
+ " — likely reaped or canceled"
847
+ )
848
+ return
849
+ if task.requeue is not None:
850
+ logger.info(
851
+ f"Task '{name}' (id={task_id}) ended as {status.value};"
852
+ " re-queued with the version upserted while it ran"
853
+ )
854
+ task.apply_requeue()
855
+ return
856
+ task.status = status
857
+ task.completed_at = datetime.now(UTC)
858
+ if error_msg:
859
+ task.error = error_msg
860
+ except Exception as e:
861
+ logger.error(f"Error updating task status: {e}")
862
+
863
+
759
864
  def _process_one_off_task(
760
865
  pq: PQ,
761
866
  *,
@@ -798,9 +903,11 @@ def _process_one_off_task(
798
903
  if task is None:
799
904
  return False
800
905
 
801
- # Mark as running
906
+ # Mark as running. ``started_at`` also identifies this run when
907
+ # Phase 3 records its result (see ``_finish_one_off``).
908
+ started_at = datetime.now(UTC)
802
909
  task.status = TaskStatus.RUNNING
803
- task.started_at = datetime.now(UTC)
910
+ task.started_at = started_at
804
911
  task.attempts += 1
805
912
 
806
913
  # Get task data for execution (before session closes)
@@ -869,27 +976,9 @@ def _process_one_off_task(
869
976
  elapsed = time.perf_counter() - start
870
977
 
871
978
  # Phase 3: Update task status (guarded — only if still RUNNING, so we
872
- # don't overwrite a reaper's FAILED verdict for an orphaned task)
873
- try:
874
- with pq.session() as session:
875
- values: dict[str, object] = {
876
- "status": status,
877
- "completed_at": datetime.now(UTC),
878
- }
879
- if error_msg:
880
- values["error"] = error_msg
881
- result = session.execute(
882
- update(Task)
883
- .where(Task.id == task_id, Task.status == TaskStatus.RUNNING)
884
- .values(**values)
885
- )
886
- if result.rowcount == 0:
887
- logger.warning(
888
- f"Task '{name}' (id={task_id}) was no longer RUNNING"
889
- " when Phase 3 tried to update — likely reaped"
890
- )
891
- except Exception as e:
892
- logger.error(f"Error updating task status: {e}")
979
+ # don't overwrite a reaper's FAILED verdict for an orphaned task; and
980
+ # re-queued instead of closed if an upsert was parked on the row)
981
+ _finish_one_off(pq, task_id, started_at, name, status, error_msg)
893
982
 
894
983
  # Log result
895
984
  if status == TaskStatus.COMPLETED:
@@ -1092,8 +1181,9 @@ def _claim_and_fork_one_off(
1092
1181
  if task is None:
1093
1182
  return None
1094
1183
 
1184
+ started_at = datetime.now(UTC)
1095
1185
  task.status = TaskStatus.RUNNING
1096
- task.started_at = datetime.now(UTC)
1186
+ task.started_at = started_at
1097
1187
  task.attempts += 1
1098
1188
 
1099
1189
  name = task.name
@@ -1128,15 +1218,7 @@ def _claim_and_fork_one_off(
1128
1218
  )
1129
1219
  except Exception as e:
1130
1220
  logger.error(f"Error starting task '{name}': {e}")
1131
- try:
1132
- with pq.session() as session:
1133
- t = session.get(Task, task_id)
1134
- if t:
1135
- t.status = TaskStatus.FAILED
1136
- t.completed_at = datetime.now(UTC)
1137
- t.error = str(e)
1138
- except Exception as update_err:
1139
- logger.error(f"Error updating task status: {update_err}")
1221
+ _finish_one_off(pq, task_id, started_at, name, TaskStatus.FAILED, str(e))
1140
1222
  return None
1141
1223
 
1142
1224
  return _ChildSlot(
@@ -1146,6 +1228,7 @@ def _claim_and_fork_one_off(
1146
1228
  name=name,
1147
1229
  start_time=time.perf_counter(),
1148
1230
  is_periodic=False,
1231
+ started_at=started_at,
1149
1232
  )
1150
1233
 
1151
1234
 
@@ -1342,29 +1425,14 @@ def _reap_and_update(pq: PQ, slot: _ChildSlot) -> None:
1342
1425
  if result.exit_kind == "timeout":
1343
1426
  error_msg = f"Timed out after {elapsed:.3f} s"
1344
1427
 
1345
- try:
1346
- with pq.session() as session:
1347
- values: dict[str, object] = {
1348
- "status": result.task_status,
1349
- "completed_at": datetime.now(UTC),
1350
- }
1351
- if error_msg:
1352
- values["error"] = error_msg
1353
- row_result = session.execute(
1354
- update(Task)
1355
- .where(
1356
- Task.id == slot.task_id,
1357
- Task.status == TaskStatus.RUNNING,
1358
- )
1359
- .values(**values)
1360
- )
1361
- if row_result.rowcount == 0:
1362
- logger.warning(
1363
- f"Task '{slot.name}' (id={slot.task_id}) was no longer"
1364
- " RUNNING when reap tried to update — likely reaped"
1365
- )
1366
- except Exception as e:
1367
- logger.error(f"Error updating task status: {e}")
1428
+ _finish_one_off(
1429
+ pq,
1430
+ slot.task_id,
1431
+ slot.started_at,
1432
+ slot.name,
1433
+ result.task_status,
1434
+ error_msg,
1435
+ )
1368
1436
 
1369
1437
  if result.task_status == TaskStatus.COMPLETED:
1370
1438
  logger.debug(f"Task '{slot.name}' completed in {elapsed:.3f} s")
@@ -1450,6 +1518,7 @@ def _run_concurrent(
1450
1518
  # drain below still runs, just without a deadline.
1451
1519
  pass
1452
1520
 
1521
+ _shutdown.log_request()
1453
1522
  if children:
1454
1523
  logger.info(f"Shutting down, draining {len(children)} in-flight task(s)...")
1455
1524
  # _wait_for_child inside _reap_and_update enforces the drain
File without changes
File without changes
File without changes
File without changes
File without changes