如何在Rust的Kafka生产者中正确启用批处理提升性能

问题描述 #

我用Rust开发了一个向Kafka Topic发送消息的应用，性能远低于预期：单条消息发送延迟约200ms，每秒仅能发送5条消息。尝试通过批处理优化，将queue.buffering.max.ms配置为500ms启用批处理后，消息发送间隔反而增至约540ms。修改rust-rdkafka的simple_producer示例复现了该问题，代码及运行输出如下：

复现代码 #

use std::time::Duration;

use log::info;

use rdkafka::config::ClientConfig;
use rdkafka::producer::{FutureProducer, FutureRecord};
use tokio::time::Instant;

#[macro_use]
extern crate env_var;

async fn produce(brokers: &str, topic_name: &str, kafka_user: &str, kafka_password: &str) {
    let producer: &FutureProducer = &ClientConfig::new()
        .set("bootstrap.servers", brokers.to_string())
        .set("security.protocol", "SASL_SSL")
        .set("sasl.mechanism", "PLAIN")
        .set("sasl.username", kafka_user.to_string())
        .set("sasl.password", kafka_password.to_string())
        .set("queue.buffering.max.ms", "500".to_string())
        .set("message.timeout.ms", "5000")
        .create()
        .expect("Producer creation error");

    // This loop is non blocking: all messages will be sent one after the other, without waiting
    // for the results.
    let futures = (0..5)
        .map(|i| async move {
            // The send operation on the topic returns a future, which will be
            // completed once the result or failure from Kafka is received.
            let delivery_status = producer
                .send(
                    FutureRecord::to(topic_name)
                        .payload(&format!("Message {}", i))
                        .key(&format!("Key {}", i)),
                    Duration::from_secs(0),
                )
                .await;

            // This will be executed when the result is received.
            info!("Delivery status for message {} received", i);
            delivery_status
        })
        .collect::<Vec<_>>();

    // This loop will wait until all delivery statuses have been received.
    
    for future in futures {
        let timer = Instant::now();
        let result = future.await;
        info!("Future completed after {:?}. Result: {:?}", timer.elapsed(), result);
    }
}

#[tokio::main]
async fn main() {
    env_logger::init();
    let kafka_servers = env_var!(required "KAFKA_SERVERS");
    let kafka_user = env_var!(optional "KAFKA_USER", default: "$ConnectionString");
    let kafka_password = env_var!(required "KAFKA_PASSWORD");
    let kafka_topic = env_var!(required "KAFKA_TOPIC");

    produce(
        kafka_servers.as_str(),
        kafka_topic.as_str(),
        kafka_user.as_str(),
        kafka_password.as_str(),
    )
    .await;
}

运行输出 #

[2022-09-26T13:53:50Z INFO  batch_test] Delivery status for message 0 received
[2022-09-26T13:53:50Z INFO  batch_test] Future completed after 922.2421ms. Result: Ok((0, 414835))
[2022-09-26T13:53:50Z INFO  batch_test] Delivery status for message 1 received
[2022-09-26T13:53:50Z INFO  batch_test] Future completed after 543.6226ms. Result: Ok((0, 414837))
[2022-09-26T13:53:51Z INFO  batch_test] Delivery status for message 2 received
[2022-09-26T13:53:51Z INFO  batch_test] Future completed after 549.2269ms. Result: Ok((0, 414838))
[2022-09-26T13:53:51Z INFO  batch_test] Delivery status for message 3 received
[2022-09-26T13:53:51Z INFO  batch_test] Future completed after 548.8964ms. Result: Ok((0, 414839))
[2022-09-26T13:53:52Z INFO  batch_test] Delivery status for message 4 received
[2022-09-26T13:53:52Z INFO  batch_test] Future completed after 538.4841ms. Result: Ok((0, 414840))

问题原因与解决方案 #

1. 核心问题：串行等待发送导致无法攒批 #

原代码的发送逻辑存在本质缺陷：创建future时，每个send调用都直接await，随后又逐个等待每个future完成。这意味着生产者必须等前一条消息发送完成并收到确认后，才会发送下一条消息，完全没有机会将多条消息攒成一批发送，queue.buffering.max.ms的配置自然无法生效。

修正后的发送逻辑

将send的future先收集起来，再用tokio::join_all同时等待所有future完成，让生产者有足够的消息来攒批：

async fn produce(brokers: &str, topic_name: &str, kafka_user: &str, kafka_password: &str) {
    let producer: FutureProducer = ClientConfig::new()
        .set("bootstrap.servers", brokers.to_string())
        .set("security.protocol", "SASL_SSL")
        .set("sasl.mechanism", "PLAIN")
        .set("sasl.username", kafka_user.to_string())
        .set("sasl.password", kafka_password.to_string())
        .set("queue.buffering.max.ms", "100") // 调整为更合理的延迟
        .set("queue.buffering.max.messages", "1000") // 单批最大消息数
        .set("batch.size", "16384") // 单批最大字节数（16KB）
        .set("message.timeout.ms", "5000")
        .create()
        .expect("Producer creation error");

    // 收集所有send的future，不提前await
    let futures = (0..5)
        .map(|i| {
            producer.send(
                FutureRecord::to(topic_name)
                    .payload(&format!("Message {}", i))
                    .key(&format!("Key {}", i)),
                Duration::from_secs(0),
            )
        })
        .collect::<Vec<_>>();

    // 同时等待所有发送完成
    let start_time = Instant::now();
    let results = tokio::join_all(futures).await;
    
    for (i, result) in results.into_iter().enumerate() {
        info!("Message {} delivery result: {:?}, total elapsed: {:?}", 
              i, result, start_time.elapsed());
    }
}

2. 优化批处理配置 #

仅设置queue.buffering.max.ms不足以高效触发批处理，需要结合其他配置让生产者在满足任一条件时就发送批次：

• queue.buffering.max.ms：消息在缓冲区的最大等待时间，到点就发送批次
• queue.buffering.max.messages：单个批次的最大消息数，达到该数量立即发送
• batch.size：单个批次的最大字节数，达到该大小立即发送
• 可以根据业务需求调整这些值，比如降低queue.buffering.max.ms到100ms，平衡吞吐量和延迟

3. 验证批处理生效 #

修改后运行代码，会发现所有消息的总耗时接近queue.buffering.max.ms的配置值，而不是每条消息都等待几百毫秒，说明批处理已经生效。

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