feat(kafka source): process messages in batch

src/sources/kafka.rs

@@ -15,13 +15,13 @@ use chrono::{DateTime, TimeZone, Utc};
 use futures::{Stream, StreamExt};
 use futures_util::future::OptionFuture;
 use rdkafka::{
-    ClientConfig, ClientContext, Statistics, TopicPartitionList,
+    ClientConfig, ClientContext, Statistics, Timestamp, TopicPartitionList,
     consumer::{
         BaseConsumer, CommitMode, Consumer, ConsumerContext, Rebalance, StreamConsumer,
         stream_consumer::StreamPartitionQueue,
     },
     error::KafkaError,
-    message::{BorrowedMessage, Headers as _, Message},
+    message::{BorrowedHeaders, BorrowedMessage, Headers as _, Message, OwnedHeaders},
     types::RDKafkaErrorCode,
 };
 use serde_with::serde_as;
@@ -44,8 +44,10 @@ use vector_lib::{
     },
     config::{LegacyKey, LogNamespace},
     configurable::configurable_component,
+    event::{BatchStatus, BatchStatusReceiver},
     finalizer::OrderedFinalizer,
     lookup::{OwnedValuePath, lookup_v2::OptionalValuePath, owned_value_path, path},
+    source_sender::CHUNK_SIZE,
 };
 use vrl::value::{Kind, ObjectMap, kind::Collection};
 
@@ -56,7 +58,7 @@ use crate::{
         LogSchema, SourceAcknowledgementsConfig, SourceConfig, SourceContext, SourceOutput,
         log_schema,
     },
-    event::{BatchNotifier, BatchStatus, Event, Value},
+    event::{BatchNotifier, Event, Value},
     internal_events::{
         KafkaBytesReceived, KafkaEventsReceived, KafkaOffsetUpdateError, KafkaReadError,
         StreamClosedError,
@@ -538,12 +540,12 @@ enum ConsumerState {
     Complete,
 }
 impl Draining {
-    fn new(signal: SyncSender<()>, shutdown: bool, span: Span) -> Self {
+    fn new(signal: SyncSender<()>, shutdown: bool, state: Consuming) -> Self {
         Self {
             signal,
             shutdown,
             expect_drain: HashSet::new(),
-            span,
+            span: state.span,
         }
     }
 
@@ -596,7 +598,7 @@ impl ConsumerStateInner<Consuming> {
         let (end_tx, mut end_signal) = oneshot::channel::<()>();
 
         let handle = join_set.spawn(async move {
-            let mut messages = p.stream();
+            let mut messages = p.stream().ready_chunks(CHUNK_SIZE);
             let (finalizer, mut ack_stream) = OrderedFinalizer::<FinalizerEntry>::new(None);
 
             // finalizer is the entry point for new pending acknowledgements;
@@ -621,7 +623,7 @@ impl ConsumerStateInner<Consuming> {
                     ack = ack_stream.next() => match ack {
                         Some((status, entry)) => {
                             if status == BatchStatus::Delivered
-                                && let Err(error) =  consumer.store_offset(&entry.topic, entry.partition, entry.offset) {
+                                && let Err(error) = consumer.store_offset(&entry.topic, entry.partition, entry.offset) {
                                     emit!(KafkaOffsetUpdateError { error });
                                 }
                         }
@@ -636,22 +638,43 @@ impl ConsumerStateInner<Consuming> {
 
                     message = messages.next(), if finalizer.is_some() => match message {
                         None => unreachable!("MessageStream never calls Ready(None)"),
-                        Some(Err(error)) => match error {
-                            rdkafka::error::KafkaError::PartitionEOF(partition) if exit_eof => {
-                                debug!("EOF for partition {}.", partition);
-                                status = PartitionConsumerStatus::PartitionEOF;
-                                finalizer.take();
-                            },
-                            _ => emit!(KafkaReadError { error }),
-                        },
-                        Some(Ok(msg)) => {
-                            emit!(KafkaBytesReceived {
-                                byte_size: msg.payload_len(),
-                                protocol: "tcp",
-                                topic: msg.topic(),
-                                partition: msg.partition(),
-                            });
-                            parse_message(msg, decoder.clone(), &keys, &mut out, acknowledgements, &finalizer, log_namespace).await;
+                        Some(msgs) => {
+                            let mut oks = Vec::default();
+                            let mut errors = Vec::default();
+                            for msg in msgs {
+                                match msg {
+                                    Ok(msg) => oks.push(msg),
+                                    Err(err) => {
+                                        if matches!(err, rdkafka::error::KafkaError::PartitionEOF(_)) {
+                                            errors.push(err);
+                                            break;
+                                        }
+                                        errors.push(err);
+                                    }
+                                }
+                            }
+
+                            // Detach messages from rdkafka early - this duplicates some memory,
+                            // but is needed for multithreading. Parsing has to copy data
+                            // anyways, so it just takes data from the detached message.
+                            let msgs = oks.into_iter().filter_map(|b|
+                                // The only case TryInto will fail if the message is empty
+                                // And we want to ignore empty messages
+                                b.try_into().ok()
+                            ).collect();
+                            let batch_result = parse_message(msgs, &decoder, &keys, &mut out, acknowledgements, log_namespace).await;
+                            Self::finalize_batch(batch_result, finalizer.as_ref());
+
+                            for error in errors  {
+                                match error {
+                                    rdkafka::error::KafkaError::PartitionEOF(partition) if exit_eof => {
+                                        debug!("EOF for partition {}.", partition);
+                                        status = PartitionConsumerStatus::PartitionEOF;
+                                        finalizer.take();
+                                    },
+                                    _ => emit!(KafkaReadError { error }),
+                                }
+                            }
                         }
                     },
                 )
@@ -676,7 +699,7 @@ impl ConsumerStateInner<Consuming> {
             decoder: self.decoder,
             out: self.out,
             log_namespace: self.log_namespace,
-            consumer_state: Draining::new(sig, shutdown, self.consumer_state.span),
+            consumer_state: Draining::new(sig, shutdown, self.consumer_state),
         };
 
         (Some(deadline).into(), draining)
@@ -685,6 +708,17 @@ impl ConsumerStateInner<Consuming> {
     pub const fn keep_consuming(self, deadline: OptionDeadline) -> (OptionDeadline, ConsumerState) {
         (deadline, ConsumerState::Consuming(self))
     }
+
+    fn finalize_batch(
+        batch: Option<(FinalizerEntry, BatchStatusReceiver)>,
+        finalizer: Option<&OrderedFinalizer<FinalizerEntry>>,
+    ) {
+        if let Some((msg, receiver)) = batch
+            && let Some(f) = finalizer.as_ref()
+        {
+            f.add(msg, receiver);
+        }
+    }
 }
 
 impl ConsumerStateInner<Draining> {
@@ -949,56 +983,64 @@ fn drive_kafka_consumer(
 }
 
 async fn parse_message(
-    msg: BorrowedMessage<'_>,
-    decoder: Decoder,
-    keys: &'_ Keys,
+    messages: Vec<OwnedMessage>,
+    decoder: &Decoder,
+    keys: &Keys,
     out: &mut SourceSender,
     acknowledgements: bool,
-    finalizer: &Option<OrderedFinalizer<FinalizerEntry>>,
     log_namespace: LogNamespace,
-) {
-    if let Some((count, stream)) = parse_stream(&msg, decoder, keys, log_namespace) {
-        let (batch, receiver) = BatchNotifier::new_with_receiver();
-        let mut stream = stream.map(|event| {
-            // All acknowledgements flow through the normal Finalizer stream so
-            // that they can be handled in one place, but are only tied to the
-            // batch when acknowledgements are enabled
-            if acknowledgements {
-                event.with_batch_notifier(&batch)
-            } else {
-                event
-            }
-        });
-        match out.send_event_stream(&mut stream).await {
-            Err(_) => {
-                emit!(StreamClosedError { count });
-            }
-            Ok(_) => {
-                // Drop stream to avoid borrowing `msg`: "[...] borrow might be used
-                // here, when `stream` is dropped and runs the destructor [...]".
-                drop(stream);
-                if let Some(f) = finalizer.as_ref() {
-                    f.add(msg.into(), receiver)
-                }
-            }
+) -> Option<(FinalizerEntry, BatchStatusReceiver)> {
+    let (batch, receiver) = BatchNotifier::new_with_receiver();
+    let last: FinalizerEntry = messages.last()?.into();
+    let size = messages.len();
+    let (count, streams) = messages
+        .into_iter()
+        .filter_map(|msg| parse_stream(msg, decoder.clone(), keys, log_namespace))
+        .fold(
+            (0usize, Vec::with_capacity(size)),
+            |(lc, mut ls), (rc, rs)| {
+                ls.push(rs);
+                (lc + rc, ls)
+            },
+        );
+    let mut batch_stream = futures::stream::iter(streams).flatten().map(|event| {
+        // All acknowledgements flow through the normal Finalizer stream so
+        // that they can be handled in one place, but are only tied to the
+        // batch when acknowledgements are enabled
+        if acknowledgements {
+            event.with_batch_notifier(&batch)
+        } else {
+            event
         }
+    });
+    match out.send_event_stream(&mut batch_stream).await {
+        Err(_) => {
+            emit!(StreamClosedError { count });
+            None
+        }
+        Ok(_) => Some((last, receiver)),
     }
 }
 
 // Turn the received message into a stream of parsed events.
 fn parse_stream<'a>(
-    msg: &BorrowedMessage<'a>,
+    msg: OwnedMessage,
     decoder: Decoder,
     keys: &'a Keys,
     log_namespace: LogNamespace,
 ) -> Option<(usize, impl Stream<Item = Event> + 'a + use<'a>)> {
-    let payload = msg.payload()?; // skip messages with empty payload
-
-    let rmsg = ReceivedMessage::from(msg);
+    let size = msg.payload.len();
+    emit!(KafkaBytesReceived {
+        byte_size: size,
+        protocol: "tcp",
+        topic: &msg.topic,
+        partition: msg.partition,
+    });
+    let rmsg = ReceivedMessage::from(&msg);
 
-    let payload = Cursor::new(Bytes::copy_from_slice(payload));
+    let payload = Cursor::new(Bytes::from_owner(msg.payload));
 
-    let mut stream = DecoderFramedRead::with_capacity(payload, decoder, msg.payload_len());
+    let mut stream = DecoderFramedRead::with_capacity(payload, decoder, size);
     let (count, _) = stream.size_hint();
     let stream = stream! {
         while let Some(result) = stream.next().await {
@@ -1062,20 +1104,21 @@ struct ReceivedMessage {
 }
 
 impl ReceivedMessage {
-    fn from(msg: &BorrowedMessage<'_>) -> Self {
+    fn from(msg: &OwnedMessage) -> Self {
         // Extract timestamp from kafka message
         let timestamp = msg
-            .timestamp()
+            .timestamp
             .to_millis()
             .and_then(|millis| Utc.timestamp_millis_opt(millis).latest());
 
         let key = msg
-            .key()
+            .key
+            .as_ref()
             .map(|key| Value::from(Bytes::from(key.to_owned())))
             .unwrap_or(Value::Null);
 
         let mut headers_map = ObjectMap::new();
-        if let Some(headers) = msg.headers() {
+        if let Some(headers) = &msg.headers {
             for header in headers.iter() {
                 if let Some(value) = header.value {
                     headers_map.insert(
@@ -1090,9 +1133,9 @@ impl ReceivedMessage {
             timestamp,
             key,
             headers: headers_map,
-            topic: msg.topic().to_string(),
-            partition: msg.partition(),
-            offset: msg.offset(),
+            topic: msg.topic.to_string(),
+            partition: msg.partition,
+            offset: msg.offset,
         }
     }
 
@@ -1175,12 +1218,12 @@ struct FinalizerEntry {
     offset: i64,
 }
 
-impl<'a> From<BorrowedMessage<'a>> for FinalizerEntry {
-    fn from(msg: BorrowedMessage<'a>) -> Self {
+impl From<&OwnedMessage> for FinalizerEntry {
+    fn from(msg: &OwnedMessage) -> Self {
         Self {
-            topic: msg.topic().into(),
-            partition: msg.partition(),
-            offset: msg.offset(),
+            topic: msg.topic.clone(),
+            partition: msg.partition,
+            offset: msg.offset,
         }
     }
 }
@@ -1531,6 +1574,40 @@ mod test {
     }
 }
 
+/// Our implementation of [rdkafka::message::OwnedMessage].
+///
+/// Needed to be able to take the payload from it, without copying it again.
+#[derive(Debug, Clone)]
+struct OwnedMessage {
+    payload: Vec<u8>,
+    key: Option<Vec<u8>>,
+    topic: String,
+    timestamp: Timestamp,
+    partition: i32,
+    offset: i64,
+    headers: Option<OwnedHeaders>,
+}
+
+impl TryFrom<BorrowedMessage<'_>> for OwnedMessage {
+    type Error = ();
+
+    fn try_from(value: BorrowedMessage<'_>) -> Result<Self, Self::Error> {
+        let payload = value.payload();
+        match payload {
+            None => Err(()),
+            Some(payload) => Ok(OwnedMessage {
+                key: value.key().map(|k| k.to_vec()),
+                payload: payload.to_vec(),
+                topic: value.topic().to_owned(),
+                timestamp: value.timestamp(),
+                partition: value.partition(),
+                offset: value.offset(),
+                headers: value.headers().map(BorrowedHeaders::detach),
+            }),
+        }
+    }
+}
+
 #[cfg(feature = "kafka-integration-tests")]
 #[cfg(test)]
 mod integration_test {