85 lines
3.6 KiB
Python
85 lines
3.6 KiB
Python
# requires: aio_pika
|
||
"""RabbitMQ 延迟取消订单：消息先进带 TTL 的等待队列，过期后经死信交换机转到处理队列。
|
||
|
||
运行：AMQP_URL=amqp://guest:guest@localhost/ python delay_cancel.py
|
||
没有设置 AMQP_URL 时打印提示并正常退出。会创建并删除名字带随机后缀的交换机和队列。
|
||
"""
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import sys
|
||
import time
|
||
import uuid
|
||
|
||
import aio_pika
|
||
|
||
DELAY_MS = 1000 # 演示用 1 秒；真实业务常见是 15~30 分钟
|
||
|
||
|
||
class OrderStore:
|
||
"""模拟数据库里的订单状态（真实项目里用条件 UPDATE 保证幂等）。"""
|
||
|
||
def __init__(self) -> None:
|
||
self.status: dict[str, str] = {}
|
||
|
||
def cancel_if_unpaid(self, order_id: str) -> bool:
|
||
# 等价于 UPDATE orders SET status='cancelled' WHERE id=? AND status='pending'
|
||
if self.status.get(order_id) == "pending":
|
||
self.status[order_id] = "cancelled"
|
||
return True
|
||
return False
|
||
|
||
|
||
async def main() -> int:
|
||
url = os.environ.get("AMQP_URL")
|
||
if not url:
|
||
print("未设置 AMQP_URL，跳过（示例：amqp://guest:guest@localhost/）")
|
||
return 0
|
||
suffix = uuid.uuid4().hex[:8]
|
||
conn = await aio_pika.connect_robust(url)
|
||
async with conn:
|
||
ch = await conn.channel()
|
||
await ch.set_qos(prefetch_count=10)
|
||
dlx = await ch.declare_exchange(f"order.dlx.{suffix}", aio_pika.ExchangeType.DIRECT, auto_delete=True)
|
||
# 处理队列：真正执行“超时取消”的消费者监听这里
|
||
work_q = await ch.declare_queue(f"order.cancel.{suffix}", durable=True, auto_delete=True)
|
||
await work_q.bind(dlx, routing_key="cancel")
|
||
# 等待队列：没有消费者，消息在这里过期后转到死信交换机
|
||
wait_q = await ch.declare_queue(f"order.wait.{suffix}", durable=True, auto_delete=True, arguments={
|
||
"x-message-ttl": DELAY_MS,
|
||
"x-dead-letter-exchange": dlx.name,
|
||
"x-dead-letter-routing-key": "cancel",
|
||
})
|
||
|
||
store = OrderStore()
|
||
store.status = {"o1": "pending", "o2": "pending"}
|
||
start = time.monotonic()
|
||
for oid in ("o1", "o2", "o1"): # o1 故意发两次，模拟重复消息
|
||
await ch.default_exchange.publish(
|
||
aio_pika.Message(json.dumps({"order_id": oid}).encode(),
|
||
delivery_mode=aio_pika.DeliveryMode.PERSISTENT, message_id=f"cancel-{oid}"),
|
||
routing_key=wait_q.name)
|
||
store.status["o2"] = "paid" # o2 在超时前支付了
|
||
|
||
results: list[tuple[str, bool]] = []
|
||
async with work_q.iterator() as it:
|
||
async for msg in it:
|
||
async with msg.process(): # 正常结束自动 ack，抛异常则 nack
|
||
oid = json.loads(msg.body)["order_id"]
|
||
results.append((oid, store.cancel_if_unpaid(oid)))
|
||
if len(results) == 3:
|
||
break
|
||
elapsed = time.monotonic() - start
|
||
print("处理结果：", results, f"耗时 {elapsed:.1f}s")
|
||
assert elapsed >= DELAY_MS / 1000 * 0.9, "消息应该在延迟之后才被处理"
|
||
assert store.status == {"o1": "cancelled", "o2": "paid"}, store.status
|
||
assert sorted(results) == [("o1", False), ("o1", True), ("o2", False)] # 重复消息不会重复取消
|
||
await wait_q.delete(if_unused=False, if_empty=False)
|
||
await work_q.delete(if_unused=False, if_empty=False)
|
||
print("延迟取消正确：未支付的被取消，已支付的不受影响，重复消息幂等")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(asyncio.run(main()))
|