From 635eed51568cee990511d5d50e814c823a8ff6b3 Mon Sep 17 00:00:00 2001 From: haubur Date: Thu, 20 Aug 2026 19:22:23 +0200 Subject: [PATCH] fix: respect max_buffer_size when merging batches --- core/sdk/src/clients/producer_sharding.rs | 159 ++++++++++++++++++++-- 1 file changed, 145 insertions(+), 14 deletions(-) diff --git a/core/sdk/src/clients/producer_sharding.rs b/core/sdk/src/clients/producer_sharding.rs index 668256f83d..c2db25b42e 100644 --- a/core/sdk/src/clients/producer_sharding.rs +++ b/core/sdk/src/clients/producer_sharding.rs @@ -110,16 +110,23 @@ impl Sizeable for ShardMessage { pub struct ShardMessageWithPermit { pub inner: ShardMessage, - _bytes_permit: Option, + bytes_permit: OwnedSemaphorePermit, } impl ShardMessageWithPermit { pub fn new(msg: ShardMessage, permit_bytes: OwnedSemaphorePermit) -> Self { Self { inner: msg, - _bytes_permit: Some(permit_bytes), + bytes_permit: permit_bytes, } } + + /// Takes over `other`'s messages together with its byte permit, so the buffered bytes stay + /// charged against the `max_buffer_size` budget until the merged batch has been written. + fn merge(&mut self, other: Self) { + self.inner.messages.extend(other.inner.messages); + self.bytes_permit.merge(other.bytes_permit); + } } pub struct Shard { @@ -212,6 +219,27 @@ impl Shard { } } + /// Drains the buffer into batches. + /// + /// If adjacent ShardMessages have the same destination (stream, topic, partition combination), + /// merge to avoid multiple sends. + /// The merge happens on both, the inner [`IggyMessage`] and the [`OwnedSemaphorePermit`]. + fn merge_batches(buffer: &mut Vec) -> Vec { + let mut merged_batches: Vec = Vec::with_capacity(buffer.len()); + // Since buffer is a mutable reference, the drain leaves the buffer intact but empty, + // such that it can be filled up again. + for msg in buffer.drain(..) { + if let Some(last) = merged_batches.last_mut() + && Self::same_destination(&last.inner, &msg.inner) + { + last.merge(msg); + continue; + } + merged_batches.push(msg); + } + merged_batches + } + async fn flush_buffer( core: &Arc, slots_permit: &Arc, @@ -223,18 +251,7 @@ impl Shard { return; } - let mut merged_batches: Vec = Vec::new(); - for msg in buffer.drain(..) { - if let Some(last) = merged_batches.last_mut() - && Self::same_destination(&last.inner, &msg.inner) - { - last.inner.messages.extend(msg.inner.messages); - continue; - } - merged_batches.push(msg); - } - - for msg in merged_batches { + for msg in Self::merge_batches(buffer) { let _slot_permit = slots_permit.acquire().await; let result = core @@ -312,6 +329,120 @@ mod tests { .unwrap() } + async fn charged_batch( + budget: &Arc, + stream: Arc, + topic: Arc, + payload_size: usize, + ) -> ShardMessageWithPermit { + let message = ShardMessage { + stream, + topic, + messages: vec![dummy_message(payload_size)], + partitioning: None, + }; + let permit = budget + .clone() + .acquire_many_owned(message.get_size_bytes().as_bytes_u32()) + .await + .unwrap(); + ShardMessageWithPermit::new(message, permit) + } + + #[tokio::test] + async fn test_merge_batches_keeps_permits_of_merged_batches_charged() { + let budget = Arc::new(Semaphore::new(10_000)); + let stream = dummy_identifier(); + let topic = dummy_identifier(); + + let mut buffer = Vec::new(); + for _ in 0..3 { + buffer.push(charged_batch(&budget, stream.clone(), topic.clone(), 10).await); + } + let charged = 10_000 - budget.available_permits(); + let merged = Shard::merge_batches(&mut buffer); + + // The original buffer should be drained. + assert!(buffer.is_empty()); + assert_eq!(merged.len(), 1); + assert_eq!(merged[0].inner.messages.len(), 3); + assert_eq!(budget.available_permits(), 10_000 - charged); + + // Dropping merged gives back to the semaphore. + drop(merged); + assert_eq!(budget.available_permits(), 10_000); + } + + #[tokio::test] + async fn test_shard_keeps_budget_charged_until_merged_batch_is_written() { + const BUDGET: usize = 10_000; + + // The two channels simulate the timing of a real write, which `flush_buffer` awaits at + // `core.send_internal`. Receiving on `write_started_rx` means the worker has entered that + // call, so the batch is on the wire and its bytes are still buffered from the producer's + // point of view. Sending on `release_write_tx` simulates that the call + // returns, the loop iteration ends and the batch drops and frees its permit. + let (write_started_tx, write_started_rx) = flume::unbounded::<()>(); + let (release_write_tx, release_write_rx) = flume::unbounded::<()>(); + + let mut mock = MockProducerCoreBackend::new(); + mock.expect_send_internal() + .times(1) + .returning(move |_, _, _, _| { + let write_started_tx = write_started_tx.clone(); + let release_write_rx = release_write_rx.clone(); + Box::pin(async move { + write_started_tx.send_async(()).await.unwrap(); + release_write_rx.recv_async().await.unwrap(); + Ok(no_confirmations()) + }) + }); + + let bb = BackgroundConfig::builder() + .batch_length(3) + .batch_size(0) + .linger_time(IggyDuration::new_from_secs(60)); + let config = Arc::new(bb.build()); + + let budget = Arc::new(Semaphore::new(BUDGET)); + let slots_permit = Arc::new(Semaphore::new(100)); + + let (_stop_tx, stop_rx) = broadcast::channel(1); + let shard = Shard::new( + Arc::new(mock), + config, + slots_permit, + flume::unbounded().0, + stop_rx, + ); + + let stream = dummy_identifier(); + let topic = dummy_identifier(); + for _ in 0..3 { + let batch = charged_batch(&budget, stream.clone(), topic.clone(), 100).await; + shard.send(batch).await.unwrap(); + } + let charged = BUDGET - budget.available_permits(); + assert!(charged > 0); + + // The three batches share a destination, so the flush merges them into a single write. + write_started_rx.recv_async().await.unwrap(); + assert_eq!( + budget.available_permits(), + BUDGET - charged, + "the merged batch must hold every permit it absorbed until the write completes" + ); + + release_write_tx.send_async(()).await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), async { + while budget.available_permits() != BUDGET { + sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("the written batch must give its permits back"); + } + #[tokio::test] async fn test_shard_flushes_by_batch_length() { let mut mock = MockProducerCoreBackend::new();