如何让Python Kafka消费者处理指定消息数后优雅停止（Airflow场景）

我有一个包含BashOperator的Airflow DAG，该Operator运行Kafka生产者生成随机数量的消息，Kafka消费者会将这些消息写入JSON文件。我希望消费者在处理指定数量的消息后优雅停止，且不会导致Airflow DAG报错。曾考虑超时方案，但更倾向于简洁的实现方式，想了解业内推荐的方案，确保消费者处理完初始消息集后，DAG能无报错执行后续步骤。

生产者代码 #

import time 
import json 
import random 
from datetime import datetime
from kafka import KafkaProducer


def generate_data() -> dict: 
    id_val = random.randint(1,1000000) 
    cur_timestamp = time.strftime("%Y-%m-%d %H:%M:%S") 
    platform_type = random.choice(['ios','android','web','mobile-web']) 
    messages = {
        'id': id_val,
        'cur_timestamp': cur_timestamp,
        'type': platform_type
    }
    return messages


def serializer(messages):
    return json.dumps(messages).encode('utf-8') 


topic = "new_topic" 
producer = KafkaProducer( 
    bootstrap_servers=['localhost:9092'],
    value_serializer=serializer,
    api_version=(0,11,5)
) 

record_cnt = random.randint(10,100)
def produce_msg():
    for _ in range(record_cnt):
        send_msg = generate_data()
        producer.send(topic, send_msg)
        print(f'Producing message {str(send_msg)}')
    producer.flush()    
    producer.close()

produce_msg() 
print('produce finished')

消费者代码 #

import json 
from kafka import KafkaConsumer

topic = "new_topic" 
json_path = '/home/airflow/clickstream.json'
consumer = KafkaConsumer(topic, bootstrap_servers="localhost:9092",
enable_auto_commit=True,auto_offset_reset='earliest')

def consumer_to_json():
    with open(json_path, 'w') as json_file:
        for send_msg in consumer:
            message_value = send_msg.value.decode("utf-8")
     
            json_data = json.loads(message_value)
            print(json_data)
            json.dump(json_data, json_file)
            json_file.write('\n')
        consumer.close()
        print ('finish')  

if __name__ == "__main__":
    consumer_to_json()

方案1：通过Airflow XCom传递消息计数（最可靠） #

利用Airflow的XCom机制，让生产者将生成的消息总数传递给消费者，消费者处理对应数量的消息后主动停止。

步骤1：修改生产者代码，输出消息计数

在生产者代码末尾添加一行，将消息总数输出到标准输出，方便Airflow捕获为XCom：

# 原有代码之后添加
print(f"RECORD_CNT:{record_cnt}")

步骤2：配置Airflow DAG任务

将生产者任务的do_xcom_push设为True，让Airflow捕获标准输出内容；消费者任务通过XCom获取计数后执行：

from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator
from datetime import datetime
from kafka import KafkaConsumer
import json

def run_consumer(**context):
    # 从XCom获取生产者输出的消息计数
    xcom_output = context['ti'].xcom_pull(task_ids='run_kafka_producer')
    record_cnt = int(xcom_output.split(':')[1])
    
    # 初始化消费者并处理指定数量的消息
    consumer = KafkaConsumer(
        "new_topic", 
        bootstrap_servers="localhost:9092",
        enable_auto_commit=True,
        auto_offset_reset='earliest'
    )
    json_path = '/home/airflow/clickstream.json'
    
    with open(json_path, 'w') as json_file:
        processed_count = 0
        for send_msg in consumer:
            if processed_count >= record_cnt:
                break
            message_value = send_msg.value.decode("utf-8")
            json_data = json.loads(message_value)
            print(json_data)
            json.dump(json_data, json_file)
            json_file.write('\n')
            processed_count += 1
        consumer.close()
        print('finish')

with DAG(
    'kafka_pipeline',
    start_date=datetime(2024, 1, 1),
    schedule_interval=None
) as dag:
    producer_task = BashOperator(
        task_id='run_kafka_producer',
        bash_command='python /path/to/your/producer.py',
        do_xcom_push=True
    )
    
    consumer_task = PythonOperator(
        task_id='run_kafka_consumer',
        python_callable=run_consumer,
        provide_context=True
    )
    
    producer_task >> consumer_task

方案2：在消息中嵌入总计数标记 #

如果不想依赖Airflow XCom，可以让生产者在消息中携带总计数信息，消费者识别后停止：

修改生产者代码，发送带总计数的控制消息

在生产普通消息前，先发送一条包含总计数的控制消息：

def produce_msg():
    # 先发送控制消息，告知总消息数
    control_msg = {"type": "control", "total_count": record_cnt}
    producer.send(topic, control_msg)
    # 发送普通消息
    for _ in range(record_cnt):
        send_msg = generate_data()
        producer.send(topic, send_msg)
        print(f'Producing message {str(send_msg)}')
    producer.flush()    
    producer.close()

修改消费者代码，识别控制消息并处理对应数量

def consumer_to_json():
    consumer = KafkaConsumer(
        topic, 
        bootstrap_servers="localhost:9092",
        enable_auto_commit=True,
        auto_offset_reset='earliest'
    )
    json_path = '/home/airflow/clickstream.json'
    total_count = None
    processed_count = 0
    
    with open(json_path, 'w') as json_file:
        for send_msg in consumer:
            message_value = send_msg.value.decode("utf-8")
            json_data = json.loads(message_value)
            
            # 识别控制消息，获取总计数
            if json_data.get('type') == 'control':
                total_count = json_data['total_count']
                continue
            
            # 处理普通消息
            if total_count is not None and processed_count >= total_count:
                break
            print(json_data)
            json.dump(json_data, json_file)
            json_file.write('\n')
            processed_count += 1
        consumer.close()
        print('finish')

方案对比 #

• XCom传递方案：适合Airflow场景，可靠性高，无需修改消息格式，避免其他消息干扰。
• 消息嵌入计数方案：不依赖Airflow特性，适用于独立运行的生产者/消费者，但需注意控制消息的唯一性和识别逻辑。

内容的提问来源于stack exchange，提问作者smert97