feat: performance improvements inspired by SAQ

.gitignore

@@ -162,3 +162,5 @@ cython_debug/
 
 # ruff
 .ruff_cache
+
+.claude

CHANGELOG.md

@@ -1,5 +1,25 @@
 # Changelog
 
+## [Unreleased]
+
+### Features
+
+* Add PostgreSQL advisory locking for preventing duplicate message processing
+* Add message states (queued, active, completed) for better tracking
+* Add group-based job coordination to prevent concurrent execution of related tasks
+* Add message TTL support for automatic cleanup of completed messages
+* Add automatic sweeping of stuck messages back to queue
+* Add connection health checks and automatic reconnection
+* Add dedicated dequeue connection for improved performance
+* Implement `FOR UPDATE SKIP LOCKED` for efficient concurrent dequeuing
+* Add configurable sweep interval and stuck message timeout
+
+### Performance Improvements
+
+* Use optimized indexes for efficient dequeuing operations
+* Implement connection pooling best practices
+* Add batch operations for cleanup tasks
+
 ## [0.2.0](https://github.com/karoo-ca/taskiq-pg/compare/v0.1.7...v0.2.0) (2025-03-06)
 
 

README.md

@@ -139,6 +139,43 @@ select convert_from(result, 'UTF8') from taskiq_results;
 - `max_retry_attempts`: Maximum number of message processing attempts.
 - `connection_kwargs`: Additional arguments for asyncpg connection.
 - `pool_kwargs`: Additional arguments for asyncpg pool creation.
+- `job_lock_keyspace`: Advisory lock keyspace for jobs (default: 1).
+- `message_ttl`: Time to live for completed messages in seconds (default: 86400).
+- `stuck_message_timeout`: Time before message is considered stuck in seconds (default: 300).
+- `enable_sweeping`: Enable automatic cleanup of stuck messages (default: True).
+- `sweep_interval`: Interval between sweep operations in seconds (default: 60).
+
+## Enhanced Features
+
+### Advisory Locking
+The broker now uses PostgreSQL advisory locks to prevent duplicate message processing. Each message gets a unique lock that is held while the message is being processed and released when acknowledged.
+
+### Message States
+Messages now have three states:
+- `queued`: Message is waiting to be processed
+- `active`: Message is currently being processed
+- `completed`: Message has been processed and acknowledged
+
+### Group-based Coordination
+You can prevent concurrent execution of related tasks by setting a `group_key` in the message labels:
+
+```python
+await my_task.kicker().with_labels(group_key="user_123").kiq()
+```
+
+### Message TTL
+Control how long completed messages are retained:
+
+```python
+# Keep completed message for 1 hour
+await my_task.kicker().with_labels(ttl=3600).kiq()
+```
+
+### Automatic Cleanup
+The broker automatically:
+- Sweeps stuck messages (messages without active locks) back to the queue
+- Cleans up expired completed messages
+- Handles connection failures with automatic reconnection
 
 ## Acknowledgements
 

UPGRADE_NOTES.md

@@ -0,0 +1,100 @@
+# Upgrade Notes: Enhanced PostgreSQL Features
+
+This document describes the breaking changes and new features added to taskiq-pg inspired by SAQ's PostgreSQL implementation.
+
+## Breaking Changes
+
+### Database Schema Changes
+
+The broker now uses an enhanced database schema with additional columns:
+
+- `status`: Tracks message state (queued, active, completed)
+- `scheduled_at`: Controls when messages become available for processing
+- `lock_key`: Used for PostgreSQL advisory locking
+- `expire_at`: Automatic cleanup timestamp
+- `group_key`: For coordinating related messages
+- `retry_count`: Tracks retry attempts
+
+**Migration Required**: If you have existing messages in your database, you'll need to either:
+1. Drop and recreate the messages table (losing existing messages)
+2. Manually add the new columns with appropriate defaults
+
+### API Changes
+
+The `AsyncpgBroker` constructor now accepts additional parameters:
+- `job_lock_keyspace`: Advisory lock keyspace (default: 1)
+- `message_ttl`: Time to live for completed messages in seconds (default: 86400)
+- `stuck_message_timeout`: Time before message is considered stuck (default: 300)
+- `enable_sweeping`: Enable automatic cleanup (default: True)
+- `sweep_interval`: Interval between sweep operations (default: 60)
+
+## New Features
+
+### 1. Advisory Locking
+- Prevents duplicate message processing using PostgreSQL advisory locks
+- Each message gets a unique lock that's held during processing
+- Locks are automatically released on acknowledgment
+
+### 2. Message States
+- `queued`: Message waiting to be processed
+- `active`: Message currently being processed
+- `completed`: Message has been acknowledged
+
+### 3. Scheduled Messages
+- Messages with a `delay` label are scheduled for future processing
+- The broker efficiently handles delayed messages without blocking
+
+### 4. Group Coordination
+- Messages with the same `group_key` won't be processed concurrently
+- Useful for ensuring sequential processing of related tasks
+
+### 5. Message TTL
+- Completed messages are automatically cleaned up after TTL expires
+- Configure per-message with the `ttl` label or globally via `message_ttl`
+
+### 6. Automatic Sweeping
+- Stuck messages (no active lock) are automatically returned to queue
+- Expired messages are cleaned up periodically
+- Configurable sweep interval and timeout
+
+### 7. Connection Resilience
+- Dedicated dequeue connection with health checks
+- Automatic reconnection on connection failures
+- Better connection pool management
+
+## Usage Examples
+
+### Group Coordination
+```python
+# These tasks won't run concurrently
+await my_task.kicker().with_labels(group_key="user_123").kiq()
+await another_task.kicker().with_labels(group_key="user_123").kiq()
+```
+
+### Message TTL
+```python
+# This message will be cleaned up after 1 hour
+await my_task.kicker().with_labels(ttl=3600).kiq()
+```
+
+### Delayed Messages
+```python
+# This message will be processed after 5 minutes
+await my_task.kicker().with_labels(delay="300").kiq()
+```
+
+## Performance Improvements
+
+- `FOR UPDATE SKIP LOCKED` for efficient concurrent dequeuing
+- Optimized indexes for message queries
+- Batch operations for cleanup tasks
+- Connection pooling best practices
+
+## Monitoring
+
+The broker logs important events:
+- Swept messages returned to queue
+- Connection health issues and reconnections
+- Expired message cleanup
+
+Monitor these logs to ensure your system is operating correctly.

example.py

@@ -4,7 +4,9 @@
 **shell 1: start a worker**
 
 ```sh
-$ taskiq worker example:broker
+$ bin/pg-up
+$ export POSTGRESQL_URL="postgresql://postgres:postgres@localhost:25432/postgres"
+$ uv run taskiq worker example:broker
 [2025-01-06 11:48:14,171][taskiq.worker][INFO   ][MainProcess] Pid of a main process: 80434
 [2025-01-06 11:48:14,171][taskiq.worker][INFO   ][MainProcess] Starting 2 worker processes.
 [2025-01-06 11:48:14,175][taskiq.process-manager][INFO   ][MainProcess] Started process worker-0 with pid 80436
@@ -30,26 +32,37 @@
 **shell 2: run the example script**
 
 ```sh
-$ python example.py
+$ export POSTGRESQL_URL="postgresql://postgres:postgres@localhost:25432/postgres"
+$ uv run example.py
 is_err=False log=None return_value='All problems are solved!' execution_time=1.0 labels={} error=None
 Save reference to fc238b66b9554315b5ca3d16bcab5a8d so you can look at result later
 ```
+
+**shell 1: stop the postgres db**
+
+```sh
+$ bin/pg-down
+```
 """  # noqa: E501
 
 import asyncio
+import os
 
 from taskiq.serializers import JSONSerializer
 
 from taskiq_pg import AsyncpgBroker, AsyncpgResultBackend
 
+dsn = os.getenv(
+    "POSTGRESQL_URL",
+    "postgres://postgres:postgres@localhost:15432/postgres",
+)
+
 asyncpg_result_backend: AsyncpgResultBackend[object] = AsyncpgResultBackend(
-    dsn="postgres://postgres:postgres@localhost:15432/postgres",
+    dsn=dsn,
     serializer=JSONSerializer(),
 )
 
-broker = AsyncpgBroker(
-    dsn="postgres://postgres:postgres@localhost:15432/postgres",
-).with_result_backend(asyncpg_result_backend)
+broker = AsyncpgBroker(dsn=dsn).with_result_backend(asyncpg_result_backend)
 
 
 @broker.task()
@@ -65,7 +78,6 @@ async def worst_task_ever() -> str:
     await asyncio.sleep(10.0)
     msg = "Borked"
     raise ValueError(msg)
-    return "borked"
 
 
 async def main() -> None:

example_enhanced.py

@@ -0,0 +1,129 @@
+"""
+Example demonstrating the enhanced features of taskiq-pg with advisory locking, group coordination, and automatic cleanup.
+
+**shell 1: start a worker**
+
+```sh
+$ bin/pg-up
+$ export POSTGRESQL_URL="postgresql://postgres:postgres@localhost:25432/postgres"
+$ uv run taskiq worker example_enhanced:broker
+[2025-01-06 11:48:14,171][taskiq.worker][INFO   ][MainProcess] Pid of a main process: 80434
+[2025-01-06 11:48:14,171][taskiq.worker][INFO   ][MainProcess] Starting 2 worker processes.
+[2025-01-06 11:48:14,175][taskiq.process-manager][INFO   ][MainProcess] Started process worker-0 with pid 80436
+[2025-01-06 11:48:14,176][taskiq.process-manager][INFO   ][MainProcess] Started process worker-1 with pid 80437
+Processing update_profile for user user_123
+Processing update_settings for user user_123
+Generating daily_summary report
+```
+
+**shell 2: run the example script**
+
+```sh
+$ export POSTGRESQL_URL="postgresql://postgres:postgres@localhost:25432/postgres"
+$ uv run example_enhanced.py
+Example 1: Group coordination
+
+Example 2: Message TTL
+Task 1 result: is_err=False log=None return_value={'user_id': 'user_123', 'action': 'update_profile', 'status': 'completed'} execution_time=2.0 labels={'group_key': 'user_123'} error=None
+Task 2 result: is_err=False log=None return_value={'user_id': 'user_123', 'action': 'update_settings', 'status': 'completed'} execution_time=2.0 labels={'group_key': 'user_123'} error=None
+Task 3 result: is_err=False log=None return_value='daily_summary report generated successfully' execution_time=1.0 labels={'ttl': 3600} error=None
+```
+
+Note: Tasks 1 and 2 have the same group_key, so they execute sequentially, not concurrently.
+
+**shell 1: stop the postgres db**
+
+```sh
+$ bin/pg-down
+```
+"""  # noqa: E501
+
+import asyncio
+import os
+
+from taskiq.serializers.json_serializer import JSONSerializer
+
+from taskiq_pg import AsyncpgBroker, AsyncpgResultBackend
+
+# Connection string - update as needed for your PostgreSQL instance
+dsn = os.getenv(
+    "POSTGRESQL_URL",
+    "postgres://postgres:postgres@localhost:15432/postgres",
+)
+
+# Initialize result backend
+asyncpg_result_backend = AsyncpgResultBackend[object](
+    dsn=dsn,
+    serializer=JSONSerializer(),
+)
+
+# Initialize broker with enhanced features
+broker = AsyncpgBroker(
+    dsn=dsn,
+    job_lock_keyspace=1,  # Advisory lock keyspace
+    message_ttl=300,  # Keep completed messages for 5 minutes
+    stuck_message_timeout=60,  # Consider message stuck after 1 minute
+    enable_sweeping=True,  # Enable automatic cleanup
+    sweep_interval=30,  # Sweep every 30 seconds
+).with_result_backend(asyncpg_result_backend)
+
+
+@broker.task()
+async def process_user_data(user_id: str, action: str) -> dict[str, str]:
+    """Process user data with group coordination."""
+    print(f"process_user_data({user_id=}, {action=})")  # noqa: T201
+    await asyncio.sleep(2)  # Simulate work
+    return {"user_id": user_id, "action": action, "status": "completed"}
+
+
+@broker.task()
+async def generate_report(report_type: str) -> str:
+    """Generate a report with TTL."""
+    print(f"generate_report({report_type=})")  # noqa: T201
+    await asyncio.sleep(1)
+    return f"{report_type} report generated successfully"
+
+
+async def main() -> None:
+    """Demonstrate enhanced features."""
+    await broker.startup()
+
+    # Example 1: Group coordination
+    # These tasks won't run concurrently because they have the same group_key
+    task1 = (
+        await process_user_data.kicker()
+        .with_labels(group_key="user_123")
+        .kiq("user_123", "update_profile")
+    )
+
+    task2 = (
+        await process_user_data.kicker()
+        .with_labels(group_key="user_123")
+        .kiq("user_123", "update_settings")
+    )
+
+    # Example 2: Message TTL
+    # This message will be automatically cleaned up after 1 hour
+    task3 = (
+        await generate_report.kicker()
+        .with_labels(
+            ttl=3600  # 1 hour
+        )
+        .kiq("daily_summary")
+    )
+
+    # Wait for results
+    result1 = await task1.wait_result(timeout=5)
+    print(f"{result1=}")  # noqa: T201
+
+    result2 = await task2.wait_result(timeout=5)
+    print(f"{result2=}")  # noqa: T201
+
+    result3 = await task3.wait_result(timeout=3)
+    print(f"{result3=}")  # noqa: T201
+
+    await broker.shutdown()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

pyproject.toml

@@ -58,6 +58,9 @@ pretty = true
 show_error_codes = true
 exclude = ["examples"]
 
+[tool.pyright]
+pythonVersion = "3.9"
+
 [tool.ruff]
 line-length = 88