fix(sdk): respect max_buffer_size when merging batches

Cargo.toml

@@ -206,10 +206,10 @@ hyper-util = { version = "0.1.20", features = ["server-auto", "service"] }
 iceberg = "0.9.1"
 iceberg-catalog-rest = "0.9.1"
 iceberg-storage-opendal = "0.9.1"
-iggy = { path = "core/sdk", version = "0.11.0-edge.5" }
-iggy-cli = { path = "core/cli", version = "0.14.0-edge.5" }
-iggy_binary_protocol = { path = "core/binary_protocol", version = "0.11.0-edge.5" }
-iggy_common = { path = "core/common", version = "0.11.0-edge.5" }
+iggy = { path = "core/sdk", version = "0.11.0-edge.6" }
+iggy-cli = { path = "core/cli", version = "0.14.0-edge.6" }
+iggy_binary_protocol = { path = "core/binary_protocol", version = "0.11.0-edge.6" }
+iggy_common = { path = "core/common", version = "0.11.0-edge.6" }
 iggy_connector_sdk = { path = "core/connectors/sdk", version = "0.4.0-edge.3" }
 indexmap = "2.14.0"
 integration = { path = "core/integration" }

core/ai/mcp/Cargo.toml

@@ -17,7 +17,7 @@
 
 [package]
 name = "iggy-mcp"
-version = "0.5.0-edge.4"
+version = "0.5.0-edge.5"
 description = "MCP Server for Iggy message streaming platform"
 edition = "2024"
 license = "Apache-2.0"

core/bench/Cargo.toml

@@ -17,7 +17,7 @@
 
 [package]
 name = "iggy-bench"
-version = "0.6.0-edge.5"
+version = "0.6.0-edge.6"
 edition = "2024"
 license = "Apache-2.0"
 repository = "https://github.com/apache/iggy"

core/binary_protocol/Cargo.toml

@@ -17,7 +17,7 @@
 
 [package]
 name = "iggy_binary_protocol"
-version = "0.11.0-edge.5"
+version = "0.11.0-edge.6"
 description = "Wire protocol types and codec for the Iggy binary protocol. Shared between server and SDK."
 edition = "2024"
 rust-version.workspace = true

core/cli/Cargo.toml

@@ -17,7 +17,7 @@
 
 [package]
 name = "iggy-cli"
-version = "0.14.0-edge.5"
+version = "0.14.0-edge.6"
 edition = "2024"
 rust-version.workspace = true
 authors = ["bartosz.ciesla@gmail.com"]

core/common/Cargo.toml

@@ -17,7 +17,7 @@
 
 [package]
 name = "iggy_common"
-version = "0.11.0-edge.5"
+version = "0.11.0-edge.6"
 description = "Iggy is the persistent message streaming platform written in Rust, supporting QUIC, TCP and HTTP transport protocols, capable of processing millions of messages per second."
 edition = "2024"
 rust-version.workspace = true

core/connectors/runtime/Cargo.toml

@@ -17,7 +17,7 @@
 
 [package]
 name = "iggy-connectors"
-version = "0.5.0-edge.5"
+version = "0.5.0-edge.6"
 description = "Connectors runtime for Iggy message streaming platform"
 edition = "2024"
 license = "Apache-2.0"

core/sdk/Cargo.toml

@@ -17,7 +17,7 @@
 
 [package]
 name = "iggy"
-version = "0.11.0-edge.5"
+version = "0.11.0-edge.6"
 description = "Iggy is the persistent message streaming platform written in Rust, supporting QUIC, TCP and HTTP transport protocols, capable of processing millions of messages per second."
 edition = "2024"
 rust-version.workspace = true

core/sdk/src/clients/producer_config.rs

@@ -105,7 +105,8 @@ pub struct BackgroundConfig {
     /// Action to apply when back-pressure limits are reached
     #[builder(default = BackpressureMode::Block)]
     pub failure_mode: BackpressureMode,
-    /// Upper bound for the **bytes held in memory** across *all* shards.
+    /// Upper bound for the **bytes buffered or in flight** across *all* shards.
+    /// Bytes remain charged until the corresponding write completes.
     /// `IggyByteSize::from(0)` ⇒ unlimited.
     #[builder(default = IggyByteSize::from(32 * MIB as u64))]
     pub max_buffer_size: IggyByteSize,

core/sdk/src/clients/producer_dispatcher.rs

@@ -20,7 +20,7 @@ use crate::clients::producer_config::{BackgroundConfig, BackpressureMode};
 use crate::clients::producer_error_callback::ErrorCtx;
 use crate::clients::producer_sharding::{Shard, ShardMessage, ShardMessageWithPermit};
 use futures::FutureExt;
-use iggy_common::{Identifier, IggyError, IggyMessage, Partitioning, Sizeable};
+use iggy_common::{Identifier, IggyByteSize, IggyError, IggyMessage, Partitioning, Sizeable};
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
 use tokio::sync::{Semaphore, broadcast};
@@ -61,13 +61,16 @@ impl ProducerDispatcher {
             tracing::debug!("error-callback worker finished");
         });
 
-        let bytes_permit = {
-            let bytes = config.max_buffer_size.as_bytes_usize();
-            if bytes == 0 { usize::MAX } else { bytes }
-        };
+        let max_buffer_size = config.max_buffer_size.as_bytes_u64();
+        assert!(
+            max_buffer_size == 0 || max_buffer_size <= Semaphore::MAX_PERMITS as u64,
+            "max_buffer_size cannot exceed {} bytes on this platform",
+            Semaphore::MAX_PERMITS
+        );
+        let bytes_permit = Arc::new(Semaphore::new(max_buffer_size as usize));
 
         let slots_permit = Arc::new(Semaphore::new(if config.max_in_flight == 0 {
-            usize::MAX
+            Semaphore::MAX_PERMITS
         } else {
             config.max_in_flight
         }));
@@ -87,7 +90,7 @@ impl ProducerDispatcher {
             shards,
             config,
             closed: AtomicBool::new(false),
-            bytes_permit: Arc::new(Semaphore::new(bytes_permit)),
+            bytes_permit,
             stop_tx,
             join_handle: handle,
         }
@@ -112,41 +115,45 @@ impl ProducerDispatcher {
         };
         let batch_bytes = shard_message.get_size_bytes();
 
-        if batch_bytes > self.config.max_buffer_size {
+        if self.config.max_buffer_size != 0 && batch_bytes > self.config.max_buffer_size {
             return Err(IggyError::BackgroundSendBufferOverflow);
         }
 
-        let permit_bytes = match self
-            .bytes_permit
-            .clone()
-            .try_acquire_many_owned(batch_bytes.as_bytes_u32())
-        {
-            Ok(perm) => perm,
-            Err(_) => match self.config.failure_mode {
-                BackpressureMode::FailImmediately => {
-                    return Err(IggyError::BackgroundSendBufferOverflow);
-                }
-                BackpressureMode::Block => self
-                    .bytes_permit
-                    .clone()
-                    .acquire_many_owned(batch_bytes.as_bytes_u32())
-                    .await
-                    .map_err(|_| IggyError::BackgroundSendError)?,
-                BackpressureMode::BlockWithTimeout(timeout_dur) => {
-                    match tokio::time::timeout(
-                        timeout_dur.get_duration(),
-                        self.bytes_permit
-                            .clone()
-                            .acquire_many_owned(batch_bytes.as_bytes_u32()),
-                    )
-                    .await
-                    {
-                        Ok(Ok(perm)) => perm,
-                        Ok(Err(_)) => return Err(IggyError::BackgroundSendError),
-                        Err(_) => return Err(IggyError::BackgroundSendTimeout),
+        let permit_count = Self::permit_count(batch_bytes)?;
+        let bytes_permit = if self.config.max_buffer_size == 0 {
+            None
+        } else {
+            let permit = match self
+                .bytes_permit
+                .clone()
+                .try_acquire_many_owned(permit_count)
+            {
+                Ok(permit) => permit,
+                Err(_) => match &self.config.failure_mode {
+                    BackpressureMode::FailImmediately => {
+                        return Err(IggyError::BackgroundSendBufferOverflow);
                     }
-                }
-            },
+                    BackpressureMode::Block => self
+                        .bytes_permit
+                        .clone()
+                        .acquire_many_owned(permit_count)
+                        .await
+                        .map_err(|_| IggyError::BackgroundSendError)?,
+                    BackpressureMode::BlockWithTimeout(timeout_duration) => {
+                        match tokio::time::timeout(
+                            timeout_duration.get_duration(),
+                            self.bytes_permit.clone().acquire_many_owned(permit_count),
+                        )
+                        .await
+                        {
+                            Ok(Ok(permit)) => permit,
+                            Ok(Err(_)) => return Err(IggyError::BackgroundSendError),
+                            Err(_) => return Err(IggyError::BackgroundSendTimeout),
+                        }
+                    }
+                },
+            };
+            Some(permit)
         };
 
         let shard_ix = self.config.sharding.pick_shard(
@@ -159,10 +166,15 @@ impl ProducerDispatcher {
         let shard = &self.shards[shard_ix];
 
         shard
-            .send(ShardMessageWithPermit::new(shard_message, permit_bytes))
+            .send(ShardMessageWithPermit::new(shard_message, bytes_permit))
             .await
     }
 
+    fn permit_count(batch_size: IggyByteSize) -> Result<u32, IggyError> {
+        u32::try_from(batch_size.as_bytes_u64())
+            .map_err(|_| IggyError::BackgroundSendBufferOverflow)
+    }
+
     /// Flushes each shard's buffer and stops its worker. Dropping the
     /// dispatcher instead of calling this silently discards any buffered,
     /// not-yet-sent messages.
@@ -174,7 +186,7 @@ impl ProducerDispatcher {
         let _ = self.stop_tx.send(());
 
         for shard in self.shards.drain(..) {
-            if let Err(e) = shard._handle.await {
+            if let Err(e) = shard.handle.await {
                 tracing::error!("shard panicked: {e:?}");
             }
         }
@@ -237,6 +249,60 @@ mod tests {
         assert!(result.is_ok());
     }
 
+    #[tokio::test]
+    async fn test_dispatch_succeeds_with_unlimited_buffer_and_in_flight_requests() {
+        let mut mock = MockProducerCoreBackend::new();
+        mock.expect_send_internal()
+            .times(1)
+            .returning(|_, _, _, _| Box::pin(async { Ok(no_confirmations()) }));
+
+        let config = BackgroundConfig::builder()
+            .max_buffer_size(0.into())
+            .max_in_flight(0)
+            .batch_length(1)
+            .build();
+        let dispatcher = ProducerDispatcher::new(Arc::new(mock), config);
+
+        assert_eq!(dispatcher.bytes_permit.available_permits(), 0);
+        dispatcher
+            .dispatch(
+                vec![dummy_message(5)],
+                dummy_identifier(),
+                dummy_identifier(),
+                None,
+            )
+            .await
+            .unwrap();
+        dispatcher.shutdown().await;
+    }
+
+    #[cfg(target_pointer_width = "64")]
+    #[tokio::test]
+    async fn test_dispatcher_supports_buffer_budget_above_u32_max() {
+        let mock = MockProducerCoreBackend::new();
+        let budget_size = u32::MAX as u64 + 1;
+        let config = BackgroundConfig::builder()
+            .max_buffer_size(budget_size.into())
+            .build();
+        let dispatcher = ProducerDispatcher::new(Arc::new(mock), config);
+
+        assert_eq!(
+            dispatcher.bytes_permit.available_permits(),
+            budget_size as usize
+        );
+        dispatcher.shutdown().await;
+    }
+
+    #[test]
+    fn test_permit_count_rejects_batch_above_u32_max() {
+        let result = ProducerDispatcher::permit_count(IggyByteSize::from(u32::MAX as u64 + 1));
+
+        assert!(matches!(
+            result,
+            Err(IggyError::BackgroundSendBufferOverflow)
+        ));
+    }
+
     #[tokio::test]
     async fn test_dispatch_fails_on_buffer_overflow_immediate() {
         let mock = MockProducerCoreBackend::new();