通过消息停止消费者中某个进程的方案

生产者代码

生产者将发送包含task_id的启动和停止任务的消息到RabbitMQ队列。

producer.py

import pika
import json
import time

# 连接到RabbitMQ服务器
connection = pika.BlockingConnection(pika.ConnectionParameters('localhost'))
channel = connection.channel()

# 声明一个队列
channel.queue_declare(queue='task_queue')


def send_task(task_id, task_type):
    task_message = {
        'task_id': task_id,
        'task_type': task_type
    }
    channel.basic_publish(
        exchange='',
        routing_key='task_queue',
        body=json.dumps(task_message),
        properties=pika.BasicProperties(
            delivery_mode=2,  # 持久化消息
        )
    )
    print(f" [x] Sent {task_message}")


# 发送启动任务消息
task_id = 1
send_task(task_id, 'START_TASK')

# 等待一些时间，让任务开始执行
time.sleep(5)

# 发送停止任务消息
send_task(task_id, 'STOP_TASK')

# 关闭连接
connection.close()

消费者代码

消费者将监听队列中的消息，并根据消息中的task_id和task_type来决定是启动还是停止指定的任务。

consumer.py

import time
import pika
import multiprocessing
import json
import os


# 定义一个函数来执行耗时任务
def do_long_running_task(task_id):
    print(f"Task {task_id} started with PID: {os.getpid()}")
    try:
        while True:
            time.sleep(1)  # 模拟耗时任务
    except KeyboardInterrupt:
        print(f"Task {task_id} stopped with PID: {os.getpid()}")


# 连接到RabbitMQ服务器
connection = pika.BlockingConnection(pika.ConnectionParameters('localhost'))
channel = connection.channel()

# 声明一个队列
channel.queue_declare(queue='task_queue')

# 存储进程的字典，以task_id为键
processes = {}


# 定义回调函数来处理接收到的消息
def callback(ch, method, properties, body):
    task_message = json.loads(body)
    task_id = task_message['task_id']
    task_type = task_message['task_type']

    if task_type == 'START_TASK':
        # 启动任务
        print(f"Received START_TASK for task_id {task_id}")
        process = multiprocessing.Process(target=do_long_running_task, args=(task_id,))
        process.start()
        processes[task_id] = process
    elif task_type == 'STOP_TASK' and task_id in processes:
        # 停止指定task_id的任务
        print(f"Received STOP_TASK for task_id {task_id}")
        processes[task_id].terminate()  # 终止进程
        processes[task_id].join()  # 等待进程终止
        del processes[task_id]  # 从字典中删除进程记录


# 使用basic_consume来监听队列
channel.basic_consume(
    queue='task_queue',
    on_message_callback=callback,
    auto_ack=True
)

print(' [*] Waiting for messages. To exit press CTRL+C')
try:
    channel.start_consuming()
except KeyboardInterrupt:
    channel.stop_consuming()
    for task_id in processes:
        processes[task_id].terminate()
        processes[task_id].join()
    print("Stopped consuming and all tasks have been terminated.")
finally:
    connection.close()

运行

shell1 启动消费者

$ python consumer.py
 [*] Waiting for messages. To exit press CTRL+C

shell2 启动生产者

$ python producer.py
 [x] Sent {'task_id': 1, 'task_type': 'START_TASK'}
 [x] Sent {'task_id': 1, 'task_type': 'STOP_TASK'}

shell1 输出

$ python consumer.py
 [*] Waiting for messages. To exit press CTRL+C
Received START_TASK for task_id 1
 [*] Waiting for messages. To exit press CTRL+C
Received STOP_TASK for task_id 1