async-task-kit

A powerful async task processing kit based on RabbitMQ with asyncio consumers.

Installation

pip install async-task-kit

Features

• RabbitMQ Client: Robust connection pooling and delay/dead-letter queue support.
• CoroutineConsumer: asyncio 多 worker 消费，ConsumeSession 保证长任务 ack 可靠。
• Extensible Processor: Easily define your task logic by inheriting TaskProcessor.
• Built-in Logger & EnvLoader: Useful utilities for production-ready applications.

Quick Start

1. Configuration (.env)

Use the built-in EnvLoader to manage your environment variables. Create a .env file:

RABBITMQ_URL=amqp://guest:guest@localhost/
TASK_IDS=demo_task

# You can also configure specific task settings using the {TASK_ID}_ prefix
DEMO_TASK_QUEUE_NAME=my_demo_queue
DEMO_TASK_CONCURRENCY=3

2. Define your Task Processor (demo_processor.py)

import logging
import asyncio
from async_task_kit import TaskProcessor

logger = logging.getLogger(__name__)

class DemoProcessor(TaskProcessor):
    async def process(self, task: dict):
        logger.info(f"Processing task: {task}")
        # 阻塞 I/O 或 CPU 密集：放到线程池，避免卡住 event loop
        # result = await asyncio.to_thread(self._sync_work, task)
        # Return any truthy value (e.g., dict, object, True) for success and pass to callback.
        # Return None or False to trigger retry.
        return {"status": "ok", "processed_data": task}

    # def _sync_work(self, task: dict):
    #     ...

    async def callback(self, task: dict, result: any):
        logger.info(f"Task completed with result: {result}")

3. Main Consumer Application (main.py)

A production-ready setup with signal handling for graceful shutdown.

import asyncio
import logging
import signal
from typing import List, Type

from demo_processor import DemoProcessor

from async_task_kit import CoroutineConsumer as Consumer
from async_task_kit import TaskProcessor, EnvLoader, setup_logger

# Initialize logger
setup_logger()
logger = logging.getLogger(__name__)

# Register your processors
TASK_REGISTRY: dict[str, Type[TaskProcessor]] = {
    "demo_task": DemoProcessor,
}

consumers: List[Consumer] = []

async def run_all_consumers(amqp_url: str, task_ids: List[str]):
    tasks = []
    for task_id in task_ids:
        if task_id not in TASK_REGISTRY:
            continue

        processor_cls = TASK_REGISTRY[task_id]
        processor = processor_cls(task_id=task_id)

        consumer = Consumer(
            amqp_url=amqp_url,
            queue_name=processor.queue_name,
            processor=processor,
            concurrency=processor.concurrency,
        )
        consumers.append(consumer)
        tasks.append(consumer.start())

        logger.info(f"🚀 启动任务 [{task_id}] | queue={processor.queue_name} | 并发={processor.concurrency}")

    await asyncio.gather(*tasks)

async def shutdown_all():
    logger.info("🛑 优雅关闭所有消费者...")
    for consumer in consumers:
        await consumer.stop()
    logger.info("✅ 所有消费者已关闭")

def handle_exit_signal(*args, **kwargs):
    asyncio.create_task(shutdown_all())

async def main():
    env = EnvLoader()
    amqp_url = env.get("RABBITMQ_URL")
    task_ids_str = env.get("TASK_IDS", "").strip()

    if not task_ids_str:
        logger.warning("⚠️ 未配置 TASK_IDS")
        return

    task_ids = [t.strip() for t in task_ids_str.split(",") if t.strip()]
    valid_tasks = [t for t in task_ids if t in TASK_REGISTRY]

    loop = asyncio.get_running_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, handle_exit_signal)

    await run_all_consumers(amqp_url, valid_tasks)

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("👋 服务已安全退出")

4. Publishing Tasks (publisher.py)

import asyncio
import logging
from async_task_kit import RabbitMQ, setup_logger

setup_logger()
logger = logging.getLogger(__name__)

async def publish():
    rmq = RabbitMQ("amqp://guest:guest@localhost/")
    await rmq.init()
    
    await rmq.push("my_demo_queue", {"message": "Hello from async-task-kit!"})
    logger.info("Task published successfully.")
    
    await rmq.close()

if __name__ == "__main__":
    asyncio.run(publish())

5. RabbitMQ Client API

RabbitMQ 提供连接池、push / pop、延迟重试队列，以及 queue_length 查询队列深度。

queue_length

from async_task_kit import RabbitMQ

rmq = RabbitMQ("amqp://guest:guest@localhost/")
await rmq.init()

length = await rmq.queue_length("my_demo_queue")
if length == RabbitMQ.QUEUE_LENGTH_UNAVAILABLE:
    # -1：连不上或队列不存在，稍后重试
    ...
elif length == 0:
    # 队列存在且为空
    ...
else:
    # 当前消息数
    ...

await rmq.close()

实现要点：

• 使用 passive=True 声明：只查询已存在的队列，不会自动创建空队列。
• 断连时会关闭旧连接池并自动重建（与 push / pop 一致）。

pop 与队列声明

pop 取消息前会 declare_queue(durable=True)（默认），参数须与 push 创建队列时一致。

queue_length 使用 passive 模式，无需传 durable。

Consumer 空队列轮询

Consumer 通过 ConsumeSession（rmq.open_consume_session()）消费：每条 worker 一条长连接 channel，pop → process → ack 在同一 channel 上完成，避免 channel 池导致长任务 ack 失败。

push / queue_length 仍走连接池。

空队列时 pop 立即返回，轮询间隔由 poll_interval 控制（默认 1.0 秒）：

consumer = CoroutineConsumer(
    ...,
    poll_interval=2.0,  # 空队列时每 2 秒 pop 一次
)

6. 并发与部署

本库只提供 CoroutineConsumer（0.1.18 起移除 ThreadConsumer / ProcessConsumer）。

阻塞或 CPU 密集任务

在 process() 里用 asyncio.to_thread，把同步阻塞逻辑丢进默认线程池，不阻塞消费 event loop：

async def process(self, task: dict):
    return await asyncio.to_thread(self._heavy_sync, task)

配合 concurrency（环境变量 {TASK_ID}_CONCURRENCY）控制同一进程内并行 worker 数。

多任务类型 / 多进程隔离

不要在一个进程里混跑互不相关的重任务。推荐：

进程级隔离交给编排层（systemd、Docker、K8s），库内不再维护 ProcessConsumer。

从 ThreadConsumer / ProcessConsumer 迁移

# 旧
# from async_task_kit import ThreadConsumer as Consumer

# 新：一律 CoroutineConsumer + asyncio.to_thread（阻塞部分）
from async_task_kit import CoroutineConsumer as Consumer

Changelog

各版本变更说明见 CHANGELOG.md。

License

MIT