From 9b6a9abdf86a5ecb782ca5d702a846060c451ff2 Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 17:10:54 +0530 Subject: [PATCH 01/24] test(binary-protocol): cover poll timeout wire format --- .../src/requests/messages/poll_messages.rs | 107 ++++++++++++++---- 1 file changed, 88 insertions(+), 19 deletions(-) diff --git a/core/binary_protocol/src/requests/messages/poll_messages.rs b/core/binary_protocol/src/requests/messages/poll_messages.rs index 4c6413cfc8..380876eabb 100644 --- a/core/binary_protocol/src/requests/messages/poll_messages.rs +++ b/core/binary_protocol/src/requests/messages/poll_messages.rs @@ -17,7 +17,7 @@ use crate::WireError; use crate::WireIdentifier; -use crate::codec::{WireDecode, WireEncode, read_u8, read_u32_le}; +use crate::codec::{WireDecode, WireEncode, read_u8, read_u32_le, read_u64_le}; use crate::primitives::consumer::WireConsumer; use crate::primitives::polling_strategy::WirePollingStrategy; use bytes::{BufMut, BytesMut}; @@ -27,7 +27,7 @@ use bytes::{BufMut, BytesMut}; /// Wire format: /// ```text /// [consumer][stream_id][topic_id][partition_flag:1][partition_id:4 LE] -/// [strategy:9][count:4 LE][auto_commit:1] +/// [strategy:9][count:4 LE][auto_commit:1][wait_timeout_us:8 LE] /// ``` /// /// `partition_id` encoding: a u8 flag (1=Some, 0=None) followed by 4 bytes @@ -41,6 +41,7 @@ pub struct PollMessagesRequest { pub strategy: WirePollingStrategy, pub count: u32, pub auto_commit: bool, + pub wait_timeout_us: u64, } const PARTITION_FLAG_SIZE: usize = 1; @@ -48,6 +49,7 @@ const PARTITION_VALUE_SIZE: usize = 4; const STRATEGY_SIZE: usize = 9; const COUNT_SIZE: usize = 4; const AUTO_COMMIT_SIZE: usize = 1; +const WAIT_TIMEOUT_SIZE: usize = 8; impl WireEncode for PollMessagesRequest { fn encoded_size(&self) -> usize { @@ -59,6 +61,7 @@ impl WireEncode for PollMessagesRequest { + STRATEGY_SIZE + COUNT_SIZE + AUTO_COMMIT_SIZE + + WAIT_TIMEOUT_SIZE } fn encode(&self, buf: &mut BytesMut) { @@ -75,6 +78,7 @@ impl WireEncode for PollMessagesRequest { self.strategy.encode(buf); buf.put_u32_le(self.count); buf.put_u8(u8::from(self.auto_commit)); + buf.put_u64_le(self.wait_timeout_us); } } @@ -104,6 +108,13 @@ impl WireDecode for PollMessagesRequest { pos += 4; let auto_commit = read_u8(buf, pos)? != 0; pos += 1; + let wait_timeout_us = if buf.len() == pos { + 0 + } else { + let wait_timeout_us = read_u64_le(buf, pos)?; + pos += WAIT_TIMEOUT_SIZE; + wait_timeout_us + }; Ok(( Self { @@ -114,6 +125,7 @@ impl WireDecode for PollMessagesRequest { strategy, count, auto_commit, + wait_timeout_us, }, pos, )) @@ -124,9 +136,8 @@ impl WireDecode for PollMessagesRequest { mod tests { use super::*; - #[test] - fn roundtrip_with_partition() { - let req = PollMessagesRequest { + fn request(wait_timeout_us: u64) -> PollMessagesRequest { + PollMessagesRequest { consumer: WireConsumer::consumer(WireIdentifier::numeric(1)), stream_id: WireIdentifier::numeric(10), topic_id: WireIdentifier::numeric(20), @@ -134,7 +145,22 @@ mod tests { strategy: WirePollingStrategy::offset(100), count: 50, auto_commit: true, - }; + wait_timeout_us, + } + } + + #[test] + fn roundtrip_with_zero_wait_timeout() { + let req = request(0); + let bytes = req.to_bytes(); + let (decoded, consumed) = PollMessagesRequest::decode(&bytes).unwrap(); + assert_eq!(consumed, bytes.len()); + assert_eq!(decoded, req); + } + + #[test] + fn roundtrip_with_non_zero_wait_timeout() { + let req = request(250_000); let bytes = req.to_bytes(); let (decoded, consumed) = PollMessagesRequest::decode(&bytes).unwrap(); assert_eq!(consumed, bytes.len()); @@ -151,6 +177,7 @@ mod tests { strategy: WirePollingStrategy::first(), count: 10, auto_commit: false, + wait_timeout_us: 0, }; let bytes = req.to_bytes(); let (decoded, consumed) = PollMessagesRequest::decode(&bytes).unwrap(); @@ -168,6 +195,7 @@ mod tests { strategy: WirePollingStrategy::offset(0), count: 1, auto_commit: false, + wait_timeout_us: 1, }; let bytes = req.to_bytes(); let (decoded, consumed) = PollMessagesRequest::decode(&bytes).unwrap(); @@ -175,6 +203,53 @@ mod tests { assert_eq!(decoded, req); } + #[test] + fn legacy_request_without_wait_timeout_decodes_as_zero() { + let req = request(0); + let bytes = req.to_bytes(); + let legacy_len = bytes.len() - WAIT_TIMEOUT_SIZE; + let (decoded, consumed) = PollMessagesRequest::decode(&bytes[..legacy_len]).unwrap(); + + assert_eq!(consumed, legacy_len); + assert_eq!(decoded.wait_timeout_us, 0); + assert_eq!(decoded, req); + } + + #[test] + fn partial_trailing_wait_timeout_returns_error() { + let req = request(0); + let bytes = req.to_bytes(); + let legacy_len = bytes.len() - WAIT_TIMEOUT_SIZE; + + for timeout_bytes in 1..WAIT_TIMEOUT_SIZE { + assert!( + PollMessagesRequest::decode(&bytes[..legacy_len + timeout_bytes]).is_err(), + "expected error with {timeout_bytes} trailing timeout bytes" + ); + } + } + + #[test] + fn encoded_size_includes_wait_timeout() { + let req = request(42_000); + let bytes = req.to_bytes(); + let legacy_size = req.consumer.encoded_size() + + req.stream_id.encoded_size() + + req.topic_id.encoded_size() + + PARTITION_FLAG_SIZE + + PARTITION_VALUE_SIZE + + STRATEGY_SIZE + + COUNT_SIZE + + AUTO_COMMIT_SIZE; + + assert_eq!(req.encoded_size(), legacy_size + WAIT_TIMEOUT_SIZE); + assert_eq!(bytes.len(), req.encoded_size()); + assert_eq!( + &bytes[legacy_size..legacy_size + WAIT_TIMEOUT_SIZE], + &req.wait_timeout_us.to_le_bytes() + ); + } + #[test] fn partition_none_encodes_zero_bytes() { let req = PollMessagesRequest { @@ -185,13 +260,13 @@ mod tests { strategy: WirePollingStrategy::first(), count: 1, auto_commit: false, + wait_timeout_us: 0, }; let bytes = req.to_bytes(); - // After consumer(7) + stream_id(6) + topic_id(6) = offset 19 let partition_offset = req.consumer.encoded_size() + req.stream_id.encoded_size() + req.topic_id.encoded_size(); - assert_eq!(bytes[partition_offset], 0); // flag = 0 + assert_eq!(bytes[partition_offset], 0); assert_eq!( &bytes[partition_offset + 1..partition_offset + 5], &[0, 0, 0, 0] @@ -199,18 +274,12 @@ mod tests { } #[test] - fn truncated_returns_error() { - let req = PollMessagesRequest { - consumer: WireConsumer::consumer(WireIdentifier::numeric(1)), - stream_id: WireIdentifier::numeric(1), - topic_id: WireIdentifier::numeric(1), - partition_id: Some(1), - strategy: WirePollingStrategy::offset(0), - count: 1, - auto_commit: false, - }; + fn truncated_required_fields_return_error() { + let req = request(0); let bytes = req.to_bytes(); - for i in 0..bytes.len() { + let legacy_len = bytes.len() - WAIT_TIMEOUT_SIZE; + + for i in 0..legacy_len { assert!( PollMessagesRequest::decode(&bytes[..i]).is_err(), "expected error for truncation at byte {i}" From 88d18ac32d1b927ec21c353ced98cf7355302847 Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 17:12:24 +0530 Subject: [PATCH 02/24] test(binary-protocol): tighten poll timeout byte assertion --- core/binary_protocol/src/requests/messages/poll_messages.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/core/binary_protocol/src/requests/messages/poll_messages.rs b/core/binary_protocol/src/requests/messages/poll_messages.rs index 380876eabb..82a6248f66 100644 --- a/core/binary_protocol/src/requests/messages/poll_messages.rs +++ b/core/binary_protocol/src/requests/messages/poll_messages.rs @@ -241,12 +241,13 @@ mod tests { + STRATEGY_SIZE + COUNT_SIZE + AUTO_COMMIT_SIZE; + let wait_timeout_bytes = req.wait_timeout_us.to_le_bytes(); assert_eq!(req.encoded_size(), legacy_size + WAIT_TIMEOUT_SIZE); assert_eq!(bytes.len(), req.encoded_size()); assert_eq!( &bytes[legacy_size..legacy_size + WAIT_TIMEOUT_SIZE], - &req.wait_timeout_us.to_le_bytes() + wait_timeout_bytes.as_slice() ); } From fb6c72daefade2b3424426b0d22010c2f953259c Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 17:36:42 +0530 Subject: [PATCH 03/24] feat(common): add timeout-capable poll API --- core/common/src/traits/message_client.rs | 32 ++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/core/common/src/traits/message_client.rs b/core/common/src/traits/message_client.rs index aa332d618c..f61a6932cb 100644 --- a/core/common/src/traits/message_client.rs +++ b/core/common/src/traits/message_client.rs @@ -19,6 +19,7 @@ use crate::{ Consumer, Identifier, IggyError, IggyMessage, Partitioning, PolledMessages, PollingStrategy, }; use async_trait::async_trait; +use std::time::Duration; /// This trait defines the methods to interact with the messaging module. #[async_trait] @@ -40,6 +41,37 @@ pub trait MessageClient { auto_commit: bool, ) -> Result; + /// Polls messages and waits up to the timeout when no messages are immediately available. + /// A zero timeout preserves immediate polling behavior. + #[allow(clippy::too_many_arguments)] + async fn poll_messages_with_timeout( + &self, + stream_id: &Identifier, + topic_id: &Identifier, + partition_id: Option, + consumer: &Consumer, + strategy: &PollingStrategy, + count: u32, + auto_commit: bool, + wait_timeout: Duration, + ) -> Result { + if wait_timeout.is_zero() { + return self + .poll_messages( + stream_id, + topic_id, + partition_id, + consumer, + strategy, + count, + auto_commit, + ) + .await; + } + + Err(IggyError::FeatureUnavailable) + } + /// Send messages using specified partitioning strategy to the given stream and topic by unique IDs or names. /// /// Authentication is required, and the permission to send the messages. From db7f20237e135abd92bd867f74589ce1ea94a321 Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 17:39:15 +0530 Subject: [PATCH 04/24] feat(common): send poll timeout over binary protocol --- .../src/traits/binary_impls/messages.rs | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/core/common/src/traits/binary_impls/messages.rs b/core/common/src/traits/binary_impls/messages.rs index 24c2a799d0..332a18f0f1 100644 --- a/core/common/src/traits/binary_impls/messages.rs +++ b/core/common/src/traits/binary_impls/messages.rs @@ -42,12 +42,20 @@ use iggy_binary_protocol::requests::messages::{ }; #[cfg(feature = "vsr")] use iggy_binary_protocol::responses::consumer_groups::SyncConsumerGroupResponse; +use std::time::Duration; /// Max attempts to resolve a fenced consumer-group poll: one re-sync after the /// coordinator rejects a stale assignment, then retry once. #[cfg(feature = "vsr")] const GROUP_POLL_MAX_ATTEMPTS: usize = 2; +fn duration_to_wait_timeout_us(wait_timeout: Duration) -> Result { + wait_timeout + .as_micros() + .try_into() + .map_err(|_| IggyError::InvalidNumberValue) +} + #[cfg(feature = "vsr")] fn group_cache_key(stream_id: &Identifier, topic_id: &Identifier, group_id: &Identifier) -> String { format!("{stream_id}|{topic_id}|{group_id}") @@ -189,6 +197,7 @@ async fn poll_group_messages( strategy: &PollingStrategy, count: u32, auto_commit: bool, + wait_timeout_us: u64, ) -> Result { let key = group_cache_key(stream_id, topic_id, &consumer.id); if !client.consumer_group_state().has_assignment(&key) { @@ -218,6 +227,7 @@ async fn poll_group_messages( strategy: polling_strategy_to_wire(strategy), count, auto_commit, + wait_timeout_us, }; match client .send_raw_with_response(POLL_MESSAGES_CODE, request.to_bytes()) @@ -264,8 +274,33 @@ impl MessageClient for B { strategy: &PollingStrategy, count: u32, auto_commit: bool, + ) -> Result { + self.poll_messages_with_timeout( + stream_id, + topic_id, + partition_id, + consumer, + strategy, + count, + auto_commit, + Duration::ZERO, + ) + .await + } + + async fn poll_messages_with_timeout( + &self, + stream_id: &Identifier, + topic_id: &Identifier, + partition_id: Option, + consumer: &Consumer, + strategy: &PollingStrategy, + count: u32, + auto_commit: bool, + wait_timeout: Duration, ) -> Result { fail_if_not_authenticated(self).await?; + let wait_timeout_us = duration_to_wait_timeout_us(wait_timeout)?; // VSR: a consumer-group poll without an explicit partition is resolved // client-side from the member's cached assignment (the broker routes // explicit partitions only). @@ -279,6 +314,7 @@ impl MessageClient for B { strategy, count, auto_commit, + wait_timeout_us, ) .await; } @@ -290,6 +326,7 @@ impl MessageClient for B { strategy: polling_strategy_to_wire(strategy), count, auto_commit, + wait_timeout_us, }; let response = self .send_raw_with_response(POLL_MESSAGES_CODE, req.to_bytes()) From ef597cdd67dfd19c1b261a4f15f24b9d2d61fd29 Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 17:40:52 +0530 Subject: [PATCH 05/24] feat(sdk): expose poll timeout on Rust client --- core/sdk/src/clients/binary_message.rs | 28 +++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/core/sdk/src/clients/binary_message.rs b/core/sdk/src/clients/binary_message.rs index 7105e63c96..f85e7d610d 100644 --- a/core/sdk/src/clients/binary_message.rs +++ b/core/sdk/src/clients/binary_message.rs @@ -23,6 +23,7 @@ use iggy_common::locking::IggyRwLockFn; use iggy_common::{ Consumer, Identifier, IggyError, IggyMessage, Partitioning, PolledMessages, PollingStrategy, }; +use std::time::Duration; #[async_trait] impl MessageClient for IggyClient { @@ -35,6 +36,30 @@ impl MessageClient for IggyClient { strategy: &PollingStrategy, count: u32, auto_commit: bool, + ) -> Result { + self.poll_messages_with_timeout( + stream_id, + topic_id, + partition_id, + consumer, + strategy, + count, + auto_commit, + Duration::ZERO, + ) + .await + } + + async fn poll_messages_with_timeout( + &self, + stream_id: &Identifier, + topic_id: &Identifier, + partition_id: Option, + consumer: &Consumer, + strategy: &PollingStrategy, + count: u32, + auto_commit: bool, + wait_timeout: Duration, ) -> Result { if count == 0 { return Err(IggyError::InvalidMessagesCount); @@ -44,7 +69,7 @@ impl MessageClient for IggyClient { .client .read() .await - .poll_messages( + .poll_messages_with_timeout( stream_id, topic_id, partition_id, @@ -52,6 +77,7 @@ impl MessageClient for IggyClient { strategy, count, auto_commit, + wait_timeout, ) .await?; From ba36db4ccf39d98143e6bbe3aab7356a7bbd716b Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 17:43:35 +0530 Subject: [PATCH 06/24] feat(sdk): route poll timeout through client wrapper --- .../client_wrappers/binary_message_client.rs | 40 ++++++++++++++++--- 1 file changed, 35 insertions(+), 5 deletions(-) diff --git a/core/sdk/src/client_wrappers/binary_message_client.rs b/core/sdk/src/client_wrappers/binary_message_client.rs index 91fb303fe3..b5245731f1 100644 --- a/core/sdk/src/client_wrappers/binary_message_client.rs +++ b/core/sdk/src/client_wrappers/binary_message_client.rs @@ -21,6 +21,7 @@ use iggy_common::MessageClient; use iggy_common::{ Consumer, Identifier, IggyError, IggyMessage, Partitioning, PolledMessages, PollingStrategy, }; +use std::time::Duration; #[async_trait] impl MessageClient for ClientWrapper { @@ -33,11 +34,35 @@ impl MessageClient for ClientWrapper { strategy: &PollingStrategy, count: u32, auto_commit: bool, + ) -> Result { + self.poll_messages_with_timeout( + stream_id, + topic_id, + partition_id, + consumer, + strategy, + count, + auto_commit, + Duration::ZERO, + ) + .await + } + + async fn poll_messages_with_timeout( + &self, + stream_id: &Identifier, + topic_id: &Identifier, + partition_id: Option, + consumer: &Consumer, + strategy: &PollingStrategy, + count: u32, + auto_commit: bool, + wait_timeout: Duration, ) -> Result { match self { ClientWrapper::Iggy(client) => { client - .poll_messages( + .poll_messages_with_timeout( stream_id, topic_id, partition_id, @@ -45,12 +70,13 @@ impl MessageClient for ClientWrapper { strategy, count, auto_commit, + wait_timeout, ) .await } ClientWrapper::Http(client) => { client - .poll_messages( + .poll_messages_with_timeout( stream_id, topic_id, partition_id, @@ -58,12 +84,13 @@ impl MessageClient for ClientWrapper { strategy, count, auto_commit, + wait_timeout, ) .await } ClientWrapper::Tcp(client) => { client - .poll_messages( + .poll_messages_with_timeout( stream_id, topic_id, partition_id, @@ -71,12 +98,13 @@ impl MessageClient for ClientWrapper { strategy, count, auto_commit, + wait_timeout, ) .await } ClientWrapper::Quic(client) => { client - .poll_messages( + .poll_messages_with_timeout( stream_id, topic_id, partition_id, @@ -84,12 +112,13 @@ impl MessageClient for ClientWrapper { strategy, count, auto_commit, + wait_timeout, ) .await } ClientWrapper::WebSocket(client) => { client - .poll_messages( + .poll_messages_with_timeout( stream_id, topic_id, partition_id, @@ -97,6 +126,7 @@ impl MessageClient for ClientWrapper { strategy, count, auto_commit, + wait_timeout, ) .await } From 7366b4efd1c789e1e2c54e39ac5913bc3b3ea9c2 Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 17:53:57 +0530 Subject: [PATCH 07/24] feat(server): pass poll timeout from binary handler --- .../src/binary/handlers/messages/poll_messages_handler.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/core/server/src/binary/handlers/messages/poll_messages_handler.rs b/core/server/src/binary/handlers/messages/poll_messages_handler.rs index 6007020b7b..c8e4d6ce32 100644 --- a/core/server/src/binary/handlers/messages/poll_messages_handler.rs +++ b/core/server/src/binary/handlers/messages/poll_messages_handler.rs @@ -25,7 +25,7 @@ use crate::streaming::session::Session; use iggy_binary_protocol::requests::messages::PollMessagesRequest; use iggy_common::IggyError; use server_common::PooledBuffer; -use std::rc::Rc; +use std::{rc::Rc, time::Duration}; use tracing::{debug, trace}; pub async fn handle_poll_messages( @@ -41,13 +41,14 @@ pub async fn handle_poll_messages( let partition_id = req.partition_id; let count = req.count; let auto_commit = req.auto_commit; + let wait_timeout = Duration::from_micros(req.wait_timeout_us); debug!( "session: {session}, command: poll_messages, stream_id: {stream_id}, topic_id: {topic_id}, partition_id: {partition_id:?}" ); shard.ensure_authenticated(session)?; - let args = PollingArgs::new(strategy, count, auto_commit); + let args = PollingArgs::with_wait_timeout(strategy, count, auto_commit, wait_timeout); let user_id = session.get_user_id(); let client_id = session.client_id; From fbeba7ee0d383ee392b016ff07eafd00d78a4fbd Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 18:00:27 +0530 Subject: [PATCH 08/24] feat(server): carry poll timeout in shard args --- core/server/src/shard/system/messages.rs | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/core/server/src/shard/system/messages.rs b/core/server/src/shard/system/messages.rs index 8770ac905c..076607c98b 100644 --- a/core/server/src/shard/system/messages.rs +++ b/core/server/src/shard/system/messages.rs @@ -31,7 +31,7 @@ use iggy_common::{ }; use server_common::PooledBuffer; use server_common::sharding::IggyNamespace; -use std::sync::atomic::Ordering; +use std::{sync::atomic::Ordering, time::Duration}; use tracing::error; impl IggyShard { @@ -682,14 +682,25 @@ pub struct PollingArgs { pub strategy: PollingStrategy, pub count: u32, pub auto_commit: bool, + pub wait_timeout: Duration, } impl PollingArgs { pub fn new(strategy: PollingStrategy, count: u32, auto_commit: bool) -> Self { + Self::with_wait_timeout(strategy, count, auto_commit, Duration::ZERO) + } + + pub fn with_wait_timeout( + strategy: PollingStrategy, + count: u32, + auto_commit: bool, + wait_timeout: Duration, + ) -> Self { Self { strategy, count, auto_commit, + wait_timeout, } } } From b43b1ab86a335c35a2084f0bd999d6777a8f09e0 Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 18:35:54 +0530 Subject: [PATCH 09/24] feat(server): add poll waiter registry --- core/server/src/shard/waiters.rs | 169 +++++++++++++++++++++++++++++++ 1 file changed, 169 insertions(+) create mode 100644 core/server/src/shard/waiters.rs diff --git a/core/server/src/shard/waiters.rs b/core/server/src/shard/waiters.rs new file mode 100644 index 0000000000..139100a5b4 --- /dev/null +++ b/core/server/src/shard/waiters.rs @@ -0,0 +1,169 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::shard::IggyShard; +use ahash::AHashMap; +use async_channel::{Receiver, Sender}; +use server_common::sharding::IggyNamespace; +use std::{cell::RefCell, time::Instant}; +use std::time::Duration; + +const MAX_WAITERS_PER_NAMESPACE: usize = 1024; + +#[derive(Debug)] +struct PollWaiter { + id: u64, + wake_sender: Sender<()>, + deadline: Option, +} + +#[derive(Debug, Default)] +pub(crate) struct PollWaiterRegistry { + next_id: u64, + waiters: AHashMap>, +} + +impl PollWaiterRegistry { + fn register( + &mut self, + namespace: IggyNamespace, + timeout: Duration, + ) -> Option<(u64, Receiver<()>)> { + self.prune_namespace(&namespace); + let waiters = self.waiters.entry(namespace).or_default(); + if waiters.len() >= MAX_WAITERS_PER_NAMESPACE { + return None; + } + + self.next_id = self.next_id.wrapping_add(1); + let id = self.next_id; + let (wake_sender, wake_receiver) = async_channel::bounded(1); + waiters.push(PollWaiter { + id, + wake_sender, + deadline: Instant::now().checked_add(timeout), + }); + + Some((id, wake_receiver)) + } + + fn remove(&mut self, namespace: &IggyNamespace, id: u64) { + let Some(waiters) = self.waiters.get_mut(namespace) else { + return; + }; + waiters.retain(|waiter| waiter.id != id); + if waiters.is_empty() { + self.waiters.remove(namespace); + } + } + + fn wake_namespace(&mut self, namespace: &IggyNamespace) { + let Some(waiters) = self.waiters.remove(namespace) else { + return; + }; + for waiter in waiters { + let _ = waiter.wake_sender.try_send(()); + } + } + + fn wake_topic(&mut self, stream_id: usize, topic_id: usize) { + let namespaces = self + .waiters + .keys() + .copied() + .filter(|namespace| { + namespace.stream_id() == stream_id && namespace.topic_id() == topic_id + }) + .collect::>(); + for namespace in namespaces { + self.wake_namespace(&namespace); + } + } + + fn wake_stream(&mut self, stream_id: usize) { + let namespaces = self + .waiters + .keys() + .copied() + .filter(|namespace| namespace.stream_id() == stream_id) + .collect::>(); + for namespace in namespaces { + self.wake_namespace(&namespace); + } + } + + fn prune_namespace(&mut self, namespace: &IggyNamespace) { + let now = Instant::now(); + let Some(waiters) = self.waiters.get_mut(namespace) else { + return; + }; + waiters.retain(|waiter| { + !waiter.wake_sender.is_closed() + && waiter.deadline.is_none_or(|deadline| deadline > now) + }); + if waiters.is_empty() { + self.waiters.remove(namespace); + } + } +} + +pub(crate) struct PollWaiterRegistration<'a> { + namespace: IggyNamespace, + id: u64, + receiver: Receiver<()>, + registry: &'a RefCell, +} + +impl PollWaiterRegistration<'_> { + pub(crate) async fn wait(&self) -> bool { + self.receiver.recv().await.is_ok() + } +} + +impl Drop for PollWaiterRegistration<'_> { + fn drop(&mut self) { + self.registry.borrow_mut().remove(&self.namespace, self.id); + } +} + +impl IggyShard { + pub(crate) fn register_poll_waiter( + &self, + namespace: IggyNamespace, + timeout: Duration, + ) -> Option> { + let (id, receiver) = self.poll_waiters.borrow_mut().register(namespace, timeout)?; + Some(PollWaiterRegistration { + namespace, + id, + receiver, + registry: &self.poll_waiters, + }) + } + + pub(crate) fn wake_poll_waiters(&self, namespace: &IggyNamespace) { + self.poll_waiters.borrow_mut().wake_namespace(namespace); + } + + pub(crate) fn wake_topic_poll_waiters(&self, stream_id: usize, topic_id: usize) { + self.poll_waiters.borrow_mut().wake_topic(stream_id, topic_id); + } + + pub(crate) fn wake_stream_poll_waiters(&self, stream_id: usize) { + self.poll_waiters.borrow_mut().wake_stream(stream_id); + } +} From 29adab3a209d01be27d5f4bbf6c8014ebf937bb1 Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 18:41:10 +0530 Subject: [PATCH 10/24] feat(server): store poll waiter registry per shard --- core/server/src/shard/mod.rs | 229 +---------------------------------- 1 file changed, 6 insertions(+), 223 deletions(-) diff --git a/core/server/src/shard/mod.rs b/core/server/src/shard/mod.rs index 03f0bad1b1..4fd6f45195 100644 --- a/core/server/src/shard/mod.rs +++ b/core/server/src/shard/mod.rs @@ -20,7 +20,10 @@ use crate::{ bootstrap::load_segments, configs::server::ServerConfig, metadata::{Metadata, MetadataWriter}, - shard::{task_registry::TaskRegistry, transmission::frame::ShardFrame}, + shard::{ + task_registry::TaskRegistry, transmission::frame::ShardFrame, + waiters::PollWaiterRegistry, + }, state::file::FileState, streaming::{ clients::client_manager::ClientManager, @@ -56,6 +59,7 @@ pub mod system; pub mod task_registry; pub mod tasks; pub mod transmission; +pub mod waiters; #[cfg(feature = "systemd")] pub mod systemd; @@ -77,6 +81,7 @@ pub struct IggyShard { pub(crate) metadata_writer: Option>, pub(crate) local_partitions: RefCell, pub(crate) pending_partition_inits: RefCell>, + pub(crate) poll_waiters: RefCell, pub(crate) shards_table: EternalPtr>, pub(crate) state: FileState, @@ -258,225 +263,3 @@ impl IggyShard { self.config .system .get_consumer_offsets_path(stream_id, topic_id, partition_id); - let consumer_group_offsets_path = self - .config - .system - .get_consumer_group_offsets_path(stream_id, topic_id, partition_id); - - // Reuse metadata's Arcs so both metadata and local_partitions - // reference the same allocation — writes via store_consumer_offset - // (metadata path) are visible to delete_oldest_segments (local path). - let consumer_offsets = init_info.consumer_offsets; - let consumer_group_offsets = init_info.consumer_group_offsets; - - { - let guard = consumer_offsets.pin(); - for co in load_consumer_offsets(&consumer_offset_path).unwrap_or_default() { - guard.insert(co.consumer_id as usize, co); - } - } - - { - let guard = consumer_group_offsets.pin(); - for (cg_id, co) in load_consumer_group_offsets(&consumer_group_offsets_path) - .unwrap_or_default() - { - guard.insert(cg_id, co); - } - } - - let message_deduplicator = - create_message_deduplicator(&self.config.system).map(Arc::new); - - match load_segments( - &self.config.system, - stream_id, - topic_id, - partition_id, - partition_path, - stats.clone(), - ) - .await - { - Ok(mut loaded_log) => { - if !loaded_log.has_segments() { - info!( - "No segments found on disk for partition ID: {} for topic ID: {} for stream ID: {}, creating initial segment", - partition_id, topic_id, stream_id - ); - let segment = crate::streaming::segments::Segment::new( - 0, - self.config.system.segment.size, - ); - let storage = - crate::streaming::segments::storage::create_segment_storage( - &self.config.system, - stream_id, - topic_id, - partition_id, - 0, - 0, - 0, - ) - .await?; - loaded_log.add_persisted_segment(segment, storage); - stats.increment_segments_count(1); - } - - // Use the max end_offset across segments that have data, - // not just the active segment. Handles the edge case where - // the active segment is empty (rotated right before shutdown). - let current_offset = loaded_log - .segments() - .iter() - .filter(|s| s.size > IggyByteSize::default()) - .map(|s| s.end_offset) - .max() - .unwrap_or(0); - stats.set_current_offset(current_offset); - - // Check if ANY segment has data. Cannot use current_offset > 0 - // because a single message at offset 0 yields current_offset = 0 - // yet must still increment on the next append. - let should_increment_offset = loaded_log - .segments() - .iter() - .any(|s| s.size > IggyByteSize::default()); - - // After a crash (OOM, SIGKILL), auto_commit may have persisted - // a consumer offset beyond what was flushed to disk. Clamp to - // the partition's actual offset to prevent permanent empty polls. - { - let guard = consumer_offsets.pin(); - for entry in guard.iter() { - let stored = entry.1.offset.load(Ordering::Relaxed); - if stored > current_offset { - warn!( - "Consumer {} offset {} ahead of partition offset {} \ - for stream {}, topic {}, partition {} - clamping \ - (crash recovery)", - entry.0, - stored, - current_offset, - stream_id, - topic_id, - partition_id - ); - entry.1.offset.store(current_offset, Ordering::Relaxed); - } - } - } - { - let guard = consumer_group_offsets.pin(); - for entry in guard.iter() { - let stored = entry.1.offset.load(Ordering::Relaxed); - if stored > current_offset { - warn!( - "Consumer group {:?} offset {} ahead of partition \ - offset {} for stream {}, topic {}, partition {} - \ - clamping (crash recovery)", - entry.0, - stored, - current_offset, - stream_id, - topic_id, - partition_id - ); - entry.1.offset.store(current_offset, Ordering::Relaxed); - } - } - } - - // Initialize journal base_offset so the three-tier - // routing in ops.rs computes correct in_memory_floor. - // Without this, journal defaults to base_offset=0 which - // causes disk reads to be skipped after restart. - if should_increment_offset { - use crate::streaming::partitions::journal::{Inner, Journal}; - loaded_log.journal_mut().init(Inner { - base_offset: current_offset + 1, - ..Default::default() - }); - } - - let revision_id = init_info.revision_id; - - let partition = LocalPartition::with_log( - loaded_log, - stats, - Arc::new(AtomicU64::new(current_offset)), - consumer_offsets, - consumer_group_offsets, - message_deduplicator, - created_at, - revision_id, - should_increment_offset, - ); - - self.local_partitions - .borrow_mut() - .insert(*namespace, partition); - - info!( - "Successfully loaded segments for stream: {}, topic: {}, partition: {}", - stream_id, topic_id, partition_id - ); - } - Err(e) => { - error!( - "Failed to load segments for stream: {}, topic: {}, partition: {}: {}", - stream_id, topic_id, partition_id, e - ); - return Err(e); - } - } - } - } - - Ok(()) - } - - async fn load_users(&self) -> Result<(), IggyError> { - let users_count = self.metadata.users_count(); - self.metrics.increment_users(users_count as u32); - info!("Initialized {} user(s).", users_count); - Ok(()) - } - - pub fn assert_init(&self) -> Result<(), IggyError> { - Ok(()) - } - - pub fn is_shutting_down(&self) -> bool { - self.is_shutting_down.load(Ordering::Relaxed) - } - - pub fn get_stop_receiver(&self) -> StopReceiver { - self.stop_receiver.clone() - } - - #[instrument(skip_all, name = "trace_shutdown")] - pub async fn trigger_shutdown(&self) -> bool { - self.is_shutting_down.store(true, Ordering::SeqCst); - debug!("Shard {} shutdown state set", self.id); - self.task_registry.graceful_shutdown(SHUTDOWN_TIMEOUT).await - } - - pub fn get_available_shards_count(&self) -> u32 { - self.shards.len() as u32 - } - - pub fn ensure_authenticated(&self, session: &Session) -> Result<(), IggyError> { - if !session.is_active() { - error!("{COMPONENT} - session is inactive, session: {session}"); - return Err(IggyError::StaleClient); - } - - if session.is_authenticated() { - Ok(()) - } else { - error!("{COMPONENT} - unauthenticated access attempt, session: {session}"); - Err(IggyError::Unauthenticated) - } - } -} From 43304d2b7a3821c97b8d226c0d79dc1c464cd55f Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 18:42:35 +0530 Subject: [PATCH 11/24] fix(server): restore shard module with poll waiters --- core/server/src/shard/mod.rs | 222 +++++++++++++++++++++++++++++++++++ 1 file changed, 222 insertions(+) diff --git a/core/server/src/shard/mod.rs b/core/server/src/shard/mod.rs index 4fd6f45195..3bc7348fcc 100644 --- a/core/server/src/shard/mod.rs +++ b/core/server/src/shard/mod.rs @@ -263,3 +263,225 @@ impl IggyShard { self.config .system .get_consumer_offsets_path(stream_id, topic_id, partition_id); + let consumer_group_offsets_path = self + .config + .system + .get_consumer_group_offsets_path(stream_id, topic_id, partition_id); + + // Reuse metadata's Arcs so both metadata and local_partitions + // reference the same allocation — writes via store_consumer_offset + // (metadata path) are visible to delete_oldest_segments (local path). + let consumer_offsets = init_info.consumer_offsets; + let consumer_group_offsets = init_info.consumer_group_offsets; + + { + let guard = consumer_offsets.pin(); + for co in load_consumer_offsets(&consumer_offset_path).unwrap_or_default() { + guard.insert(co.consumer_id as usize, co); + } + } + + { + let guard = consumer_group_offsets.pin(); + for (cg_id, co) in load_consumer_group_offsets(&consumer_group_offsets_path) + .unwrap_or_default() + { + guard.insert(cg_id, co); + } + } + + let message_deduplicator = + create_message_deduplicator(&self.config.system).map(Arc::new); + + match load_segments( + &self.config.system, + stream_id, + topic_id, + partition_id, + partition_path, + stats.clone(), + ) + .await + { + Ok(mut loaded_log) => { + if !loaded_log.has_segments() { + info!( + "No segments found on disk for partition ID: {} for topic ID: {} for stream ID: {}, creating initial segment", + partition_id, topic_id, stream_id + ); + let segment = crate::streaming::segments::Segment::new( + 0, + self.config.system.segment.size, + ); + let storage = + crate::streaming::segments::storage::create_segment_storage( + &self.config.system, + stream_id, + topic_id, + partition_id, + 0, + 0, + 0, + ) + .await?; + loaded_log.add_persisted_segment(segment, storage); + stats.increment_segments_count(1); + } + + // Use the max end_offset across segments that have data, + // not just the active segment. Handles the edge case where + // the active segment is empty (rotated right before shutdown). + let current_offset = loaded_log + .segments() + .iter() + .filter(|s| s.size > IggyByteSize::default()) + .map(|s| s.end_offset) + .max() + .unwrap_or(0); + stats.set_current_offset(current_offset); + + // Check if ANY segment has data. Cannot use current_offset > 0 + // because a single message at offset 0 yields current_offset = 0 + // yet must still increment on the next append. + let should_increment_offset = loaded_log + .segments() + .iter() + .any(|s| s.size > IggyByteSize::default()); + + // After a crash (OOM, SIGKILL), auto_commit may have persisted + // a consumer offset beyond what was flushed to disk. Clamp to + // the partition's actual offset to prevent permanent empty polls. + { + let guard = consumer_offsets.pin(); + for entry in guard.iter() { + let stored = entry.1.offset.load(Ordering::Relaxed); + if stored > current_offset { + warn!( + "Consumer {} offset {} ahead of partition offset {} \ + for stream {}, topic {}, partition {} - clamping \ + (crash recovery)", + entry.0, + stored, + current_offset, + stream_id, + topic_id, + partition_id + ); + entry.1.offset.store(current_offset, Ordering::Relaxed); + } + } + } + { + let guard = consumer_group_offsets.pin(); + for entry in guard.iter() { + let stored = entry.1.offset.load(Ordering::Relaxed); + if stored > current_offset { + warn!( + "Consumer group {:?} offset {} ahead of partition \ + offset {} for stream {}, topic {}, partition {} - \ + clamping (crash recovery)", + entry.0, + stored, + current_offset, + stream_id, + topic_id, + partition_id + ); + entry.1.offset.store(current_offset, Ordering::Relaxed); + } + } + } + + // Initialize journal base_offset so the three-tier + // routing in ops.rs computes correct in_memory_floor. + // Without this, journal defaults to base_offset=0 which + // causes disk reads to be skipped after restart. + if should_increment_offset { + use crate::streaming::partitions::journal::{Inner, Journal}; + loaded_log.journal_mut().init(Inner { + base_offset: current_offset + 1, + ..Default::default() + }); + } + + let revision_id = init_info.revision_id; + + let partition = LocalPartition::with_log( + loaded_log, + stats, + Arc::new(AtomicU64::new(current_offset)), + consumer_offsets, + consumer_group_offsets, + message_deduplicator, + created_at, + revision_id, + should_increment_offset, + ); + + self.local_partitions + .borrow_mut() + .insert(*namespace, partition); + + info!( + "Successfully loaded segments for stream: {}, topic: {}, partition: {}", + stream_id, topic_id, partition_id + ); + } + Err(e) => { + error!( + "Failed to load segments for stream: {}, topic: {}, partition: {}: {}", + stream_id, topic_id, partition_id, e + ); + return Err(e); + } + } + } + } + + Ok(()) + } + + async fn load_users(&self) -> Result<(), IggyError> { + let users_count = self.metadata.users_count(); + self.metrics.increment_users(users_count as u32); + info!("Initialized {} user(s).", users_count); + Ok(()) + } + + pub fn assert_init(&self) -> Result<(), IggyError> { + Ok(()) + } + + pub fn is_shutting_down(&self) -> bool { + self.is_shutting_down.load(Ordering::Relaxed) + } + + pub fn get_stop_receiver(&self) -> StopReceiver { + self.stop_receiver.clone() + } + + #[instrument(skip_all, name = "trace_shutdown")] + pub async fn trigger_shutdown(&self) -> bool { + self.is_shutting_down.store(true, Ordering::SeqCst); + debug!("Shard {} shutdown state set", self.id); + self.task_registry.graceful_shutdown(SHUTDOWN_TIMEOUT).await + } + + pub fn get_available_shards_count(&self) -> u32 { + self.shards.len() as u32 + } + + pub fn ensure_authenticated(&self, session: &Session) -> Result<(), IggyError> { + if !session.is_active() { + error!("{COMPONENT} - session is inactive, session: {session}"); + return Err(IggyError::StaleClient); + } + + if session.is_authenticated() { + Ok(()) + } else { + error!("{COMPONENT} - unauthenticated access attempt, session: {session}"); + Err(IggyError::Unauthenticated) + } + } +} From cec70a582b2a1b3514194f3370cb1f4834a792f8 Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 18:43:18 +0530 Subject: [PATCH 12/24] feat(server): initialize poll waiter registry --- core/server/src/shard/builder.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/core/server/src/shard/builder.rs b/core/server/src/shard/builder.rs index 817276484e..d8dc25006d 100644 --- a/core/server/src/shard/builder.rs +++ b/core/server/src/shard/builder.rs @@ -17,7 +17,7 @@ use super::{ IggyShard, TaskRegistry, transmission::connector::ShardConnector, - transmission::frame::ShardFrame, + transmission::frame::ShardFrame, waiters::PollWaiterRegistry, }; use crate::metadata::{Metadata, MetadataWriter}; use crate::streaming::partitions::local_partitions::LocalPartitions; @@ -174,6 +174,7 @@ impl IggyShardBuilder { metadata_writer: self.metadata_writer.map(RefCell::new), local_partitions, pending_partition_inits: RefCell::new(AHashSet::new()), + poll_waiters: RefCell::new(PollWaiterRegistry::default()), encryptor, config, _version: version, From e201bbe865ef81db8acaf90cd3654f6806a0528f Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 18:45:00 +0530 Subject: [PATCH 13/24] feat(server): park empty timed binary polls --- .../messages/poll_messages_handler.rs | 54 +++++++++++++++++-- 1 file changed, 50 insertions(+), 4 deletions(-) diff --git a/core/server/src/binary/handlers/messages/poll_messages_handler.rs b/core/server/src/binary/handlers/messages/poll_messages_handler.rs index c8e4d6ce32..7b9dd799d0 100644 --- a/core/server/src/binary/handlers/messages/poll_messages_handler.rs +++ b/core/server/src/binary/handlers/messages/poll_messages_handler.rs @@ -25,6 +25,7 @@ use crate::streaming::session::Session; use iggy_binary_protocol::requests::messages::PollMessagesRequest; use iggy_common::IggyError; use server_common::PooledBuffer; +use server_common::sharding::IggyNamespace; use std::{rc::Rc, time::Duration}; use tracing::{debug, trace}; @@ -48,15 +49,60 @@ pub async fn handle_poll_messages( ); shard.ensure_authenticated(session)?; - let args = PollingArgs::with_wait_timeout(strategy, count, auto_commit, wait_timeout); - let user_id = session.get_user_id(); let client_id = session.client_id; let topic = shard.resolve_topic_for_poll(user_id, &stream_id, &topic_id)?; - let (metadata, mut batch) = shard - .poll_messages(client_id, topic, consumer, partition_id, args) + let (mut metadata, mut batch) = shard + .poll_messages( + client_id, + topic, + consumer.clone(), + partition_id, + PollingArgs::with_wait_timeout(strategy.clone(), count, auto_commit, wait_timeout), + ) .await?; + if !wait_timeout.is_zero() && batch.is_empty() { + let namespace = IggyNamespace::new( + topic.stream_id, + topic.topic_id, + metadata.partition_id as usize, + ); + + if let Some(waiter) = shard.register_poll_waiter(namespace, wait_timeout) { + let immediate_args = || PollingArgs::new(strategy.clone(), count, auto_commit); + (metadata, batch) = shard + .poll_messages( + client_id, + topic, + consumer.clone(), + partition_id, + immediate_args(), + ) + .await?; + + if batch.is_empty() { + let woke = matches!( + compio::time::timeout(wait_timeout, waiter.wait()).await, + Ok(true) + ); + drop(waiter); + + if woke { + (metadata, batch) = shard + .poll_messages( + client_id, + topic, + consumer, + partition_id, + immediate_args(), + ) + .await?; + } + } + } + } + let response_length = 4 + 8 + 4 + batch.size(); let response_length_bytes = response_length.to_le_bytes(); From dbbb72aaa1d02cdc02fa805ea37b327612ebee2c Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 18:50:17 +0530 Subject: [PATCH 14/24] feat(server): wake poll waiters on partition changes --- core/server/src/shard/handlers.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/core/server/src/shard/handlers.rs b/core/server/src/shard/handlers.rs index 44cc1c0cc2..52348147dd 100644 --- a/core/server/src/shard/handlers.rs +++ b/core/server/src/shard/handlers.rs @@ -74,6 +74,7 @@ async fn handle_request( .await?; shard.metrics.increment_messages(messages_count as u64); + shard.wake_poll_waiters(&namespace); Ok(ShardResponse::SendMessages) } ShardRequestPayload::PollMessages { args, consumer } => { @@ -419,6 +420,7 @@ async fn handle_request( .await?; shard.metrics.increment_messages(messages_count as u64); + shard.wake_poll_waiters(&ns); sender.send_empty_ok_response().await?; @@ -507,12 +509,14 @@ pub async fn handle_event(shard: &Rc, event: ShardEvent) -> Result<() let mut partitions = shard.local_partitions.borrow_mut(); for partition_id in partition_ids { let ns = IggyNamespace::new(numeric_stream_id, numeric_topic_id, partition_id); + shard.wake_poll_waiters(&ns); partitions.remove(&ns); } Ok(()) } ShardEvent::PurgedStream { stream_id } => { let stream = shard.resolve_stream(&stream_id)?; + shard.wake_stream_poll_waiters(stream.id()); shard.purge_stream_local(stream).await?; Ok(()) } @@ -521,6 +525,7 @@ pub async fn handle_event(shard: &Rc, event: ShardEvent) -> Result<() topic_id, } => { let topic = shard.resolve_topic(&stream_id, &topic_id)?; + shard.wake_topic_poll_waiters(topic.stream_id, topic.topic_id); shard.purge_topic_local(topic).await?; Ok(()) } From 25db4ea192646eb287c2dd72ca35dbb7a312b2be Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 18:54:39 +0530 Subject: [PATCH 15/24] Fix poll waiter registry cleanup borrows --- core/server/src/shard/waiters.rs | 31 ++++++++++++++++++------------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/core/server/src/shard/waiters.rs b/core/server/src/shard/waiters.rs index 139100a5b4..94319f2dc7 100644 --- a/core/server/src/shard/waiters.rs +++ b/core/server/src/shard/waiters.rs @@ -19,8 +19,7 @@ use crate::shard::IggyShard; use ahash::AHashMap; use async_channel::{Receiver, Sender}; use server_common::sharding::IggyNamespace; -use std::{cell::RefCell, time::Instant}; -use std::time::Duration; +use std::{cell::RefCell, time::Duration, time::Instant}; const MAX_WAITERS_PER_NAMESPACE: usize = 1024; @@ -62,11 +61,14 @@ impl PollWaiterRegistry { } fn remove(&mut self, namespace: &IggyNamespace, id: u64) { - let Some(waiters) = self.waiters.get_mut(namespace) else { - return; + let should_remove = if let Some(waiters) = self.waiters.get_mut(namespace) { + waiters.retain(|waiter| waiter.id != id); + waiters.is_empty() + } else { + false }; - waiters.retain(|waiter| waiter.id != id); - if waiters.is_empty() { + + if should_remove { self.waiters.remove(namespace); } } @@ -108,14 +110,17 @@ impl PollWaiterRegistry { fn prune_namespace(&mut self, namespace: &IggyNamespace) { let now = Instant::now(); - let Some(waiters) = self.waiters.get_mut(namespace) else { - return; + let should_remove = if let Some(waiters) = self.waiters.get_mut(namespace) { + waiters.retain(|waiter| { + !waiter.wake_sender.is_closed() + && waiter.deadline.is_none_or(|deadline| deadline > now) + }); + waiters.is_empty() + } else { + false }; - waiters.retain(|waiter| { - !waiter.wake_sender.is_closed() - && waiter.deadline.is_none_or(|deadline| deadline > now) - }); - if waiters.is_empty() { + + if should_remove { self.waiters.remove(namespace); } } From 92743960b3cf691c4c7f1979ee1fcf899cf492e2 Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Thu, 2 Jul 2026 19:30:54 +0530 Subject: [PATCH 16/24] style: format poll waiter changes --- .../binary/handlers/messages/poll_messages_handler.rs | 8 +------- core/server/src/shard/mod.rs | 3 +-- core/server/src/shard/waiters.rs | 9 +++++++-- 3 files changed, 9 insertions(+), 11 deletions(-) diff --git a/core/server/src/binary/handlers/messages/poll_messages_handler.rs b/core/server/src/binary/handlers/messages/poll_messages_handler.rs index 7b9dd799d0..2ac28ecfed 100644 --- a/core/server/src/binary/handlers/messages/poll_messages_handler.rs +++ b/core/server/src/binary/handlers/messages/poll_messages_handler.rs @@ -90,13 +90,7 @@ pub async fn handle_poll_messages( if woke { (metadata, batch) = shard - .poll_messages( - client_id, - topic, - consumer, - partition_id, - immediate_args(), - ) + .poll_messages(client_id, topic, consumer, partition_id, immediate_args()) .await?; } } diff --git a/core/server/src/shard/mod.rs b/core/server/src/shard/mod.rs index 3bc7348fcc..66a35c4df1 100644 --- a/core/server/src/shard/mod.rs +++ b/core/server/src/shard/mod.rs @@ -21,8 +21,7 @@ use crate::{ configs::server::ServerConfig, metadata::{Metadata, MetadataWriter}, shard::{ - task_registry::TaskRegistry, transmission::frame::ShardFrame, - waiters::PollWaiterRegistry, + task_registry::TaskRegistry, transmission::frame::ShardFrame, waiters::PollWaiterRegistry, }, state::file::FileState, streaming::{ diff --git a/core/server/src/shard/waiters.rs b/core/server/src/shard/waiters.rs index 94319f2dc7..3ee68fd540 100644 --- a/core/server/src/shard/waiters.rs +++ b/core/server/src/shard/waiters.rs @@ -151,7 +151,10 @@ impl IggyShard { namespace: IggyNamespace, timeout: Duration, ) -> Option> { - let (id, receiver) = self.poll_waiters.borrow_mut().register(namespace, timeout)?; + let (id, receiver) = self + .poll_waiters + .borrow_mut() + .register(namespace, timeout)?; Some(PollWaiterRegistration { namespace, id, @@ -165,7 +168,9 @@ impl IggyShard { } pub(crate) fn wake_topic_poll_waiters(&self, stream_id: usize, topic_id: usize) { - self.poll_waiters.borrow_mut().wake_topic(stream_id, topic_id); + self.poll_waiters + .borrow_mut() + .wake_topic(stream_id, topic_id); } pub(crate) fn wake_stream_poll_waiters(&self, stream_id: usize) { From 91b30d11a7c5cee3732331abd55cca72dc9e9a6c Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Sat, 4 Jul 2026 01:04:57 +0530 Subject: [PATCH 17/24] Add deferred PollMessages wait timeout --- .../src/actors/consumer/client/interface.rs | 1 + .../src/actors/consumer/client/low_level.rs | 3 +- .../consumer/typed_benchmark_consumer.rs | 2 + .../typed_benchmark_producing_consumer.rs | 2 + core/bench/src/args/common.rs | 74 ++- core/bench/src/args/defaults.rs | 1 + core/bench/src/benchmarks/benchmark.rs | 7 +- core/bench/src/benchmarks/common.rs | 6 + core/bench/src/utils/mod.rs | 18 +- .../src/traits/binary_impls/messages.rs | 1 + core/integration/tests/server/cg.rs | 10 + .../integration/tests/server/scenarios/mod.rs | 1 + .../poll_messages_wait_timeout_scenario.rs | 429 ++++++++++++++++++ core/integration/tests/server/specific.rs | 20 + .../messages/poll_messages_handler.rs | 113 +++-- core/server/src/main.rs | 8 +- core/server/src/shard/builder.rs | 12 +- core/server/src/shard/mod.rs | 4 +- core/server/src/shard/system/utils.rs | 55 +++ core/server/src/shard/waiters.rs | 65 ++- 20 files changed, 777 insertions(+), 55 deletions(-) create mode 100644 core/integration/tests/server/scenarios/poll_messages_wait_timeout_scenario.rs diff --git a/core/bench/src/actors/consumer/client/interface.rs b/core/bench/src/actors/consumer/client/interface.rs index 5b5cc3fddc..1c52cfc190 100644 --- a/core/bench/src/actors/consumer/client/interface.rs +++ b/core/bench/src/actors/consumer/client/interface.rs @@ -27,6 +27,7 @@ pub struct BenchmarkConsumerConfig { pub stream_id: String, pub messages_per_batch: BenchmarkNumericParameter, pub warmup_time: IggyDuration, + pub poll_wait_timeout: IggyDuration, pub polling_kind: PollingKind, pub origin_timestamp_latency_calculation: bool, pub pretty: bool, diff --git a/core/bench/src/actors/consumer/client/low_level.rs b/core/bench/src/actors/consumer/client/low_level.rs index fc8c1bc644..02bf1d9be4 100644 --- a/core/bench/src/actors/consumer/client/low_level.rs +++ b/core/bench/src/actors/consumer/client/low_level.rs @@ -64,7 +64,7 @@ impl ConsumerClient for LowLevelConsumerClient { let before_poll = Instant::now(); let polled = client - .poll_messages( + .poll_messages_with_timeout( &self.stream_id, &self.topic_id, self.partition_id, @@ -72,6 +72,7 @@ impl ConsumerClient for LowLevelConsumerClient { &self.polling_strategy, messages_to_receive, self.auto_commit, + self.config.poll_wait_timeout.get_duration(), ) .await; diff --git a/core/bench/src/actors/consumer/typed_benchmark_consumer.rs b/core/bench/src/actors/consumer/typed_benchmark_consumer.rs index ae051fceda..85b9e572a3 100644 --- a/core/bench/src/actors/consumer/typed_benchmark_consumer.rs +++ b/core/bench/src/actors/consumer/typed_benchmark_consumer.rs @@ -51,6 +51,7 @@ impl TypedBenchmarkConsumer { messages_per_batch: BenchmarkNumericParameter, finish_condition: Arc, warmup_time: IggyDuration, + poll_wait_timeout: IggyDuration, sampling_time: IggyDuration, moving_average_window: u32, polling_kind: PollingKind, @@ -64,6 +65,7 @@ impl TypedBenchmarkConsumer { stream_id, messages_per_batch, warmup_time, + poll_wait_timeout, polling_kind, origin_timestamp_latency_calculation, pretty, diff --git a/core/bench/src/actors/producing_consumer/typed_benchmark_producing_consumer.rs b/core/bench/src/actors/producing_consumer/typed_benchmark_producing_consumer.rs index b193905f1f..7ba7dc66e7 100644 --- a/core/bench/src/actors/producing_consumer/typed_benchmark_producing_consumer.rs +++ b/core/bench/src/actors/producing_consumer/typed_benchmark_producing_consumer.rs @@ -58,6 +58,7 @@ impl TypedBenchmarkProducingConsumer { send_finish_condition: Arc, poll_finish_condition: Arc, warmup_time: IggyDuration, + poll_wait_timeout: IggyDuration, sampling_time: IggyDuration, moving_average_window: u32, limit_bytes_per_second: Option, @@ -81,6 +82,7 @@ impl TypedBenchmarkProducingConsumer { stream_id, messages_per_batch, warmup_time, + poll_wait_timeout, polling_kind, origin_timestamp_latency_calculation, pretty, diff --git a/core/bench/src/args/common.rs b/core/bench/src/args/common.rs index 1a6a3a46b3..8e77a34ce9 100644 --- a/core/bench/src/args/common.rs +++ b/core/bench/src/args/common.rs @@ -21,7 +21,8 @@ use super::props::{BenchmarkKindProps, BenchmarkTransportProps}; use super::{ defaults::{ DEFAULT_MESSAGE_BATCHES, DEFAULT_MESSAGE_SIZE, DEFAULT_MESSAGES_PER_BATCH, - DEFAULT_MOVING_AVERAGE_WINDOW, DEFAULT_SAMPLING_TIME, DEFAULT_WARMUP_TIME, + DEFAULT_MOVING_AVERAGE_WINDOW, DEFAULT_POLL_WAIT_TIMEOUT, DEFAULT_SAMPLING_TIME, + DEFAULT_WARMUP_TIME, }, transport::BenchmarkTransportCommand, }; @@ -71,6 +72,10 @@ pub struct IggyBenchArgs { #[arg(long, short = 'w', default_value_t = IggyDuration::from_str(DEFAULT_WARMUP_TIME).unwrap())] pub warmup_time: IggyDuration, + /// Poll wait timeout in human readable format, e.g. "10ms", "1s" + #[arg(long, default_value_t = IggyDuration::from_str(DEFAULT_POLL_WAIT_TIMEOUT).unwrap(), value_parser = IggyDuration::from_str)] + pub poll_wait_timeout: IggyDuration, + /// Sampling time for metrics collection. It is also used as bucket size for time series calculations. #[arg(long, short = 't', default_value_t = IggyDuration::from_str(DEFAULT_SAMPLING_TIME).unwrap(), value_parser = IggyDuration::from_str)] pub sampling_time: IggyDuration, @@ -163,6 +168,38 @@ impl IggyBenchArgs { .exit(); } + if self.high_level_api && !self.poll_wait_timeout.is_zero() { + Self::command() + .error( + ErrorKind::ArgumentConflict, + "--poll-wait-timeout is only supported by low-level benchmark consumers", + ) + .exit(); + } + + if matches!( + self.kind(), + BenchmarkKind::PinnedProducer | BenchmarkKind::BalancedProducer + ) && !self.poll_wait_timeout.is_zero() + { + Self::command() + .error( + ErrorKind::ArgumentConflict, + "--poll-wait-timeout requires a benchmark with consumers", + ) + .exit(); + } + + if matches!(self.transport(), TransportProtocol::Http) && !self.poll_wait_timeout.is_zero() + { + Self::command() + .error( + ErrorKind::ArgumentConflict, + "--poll-wait-timeout is not supported by HTTP transport", + ) + .exit(); + } + self.benchmark_kind.inner().validate(); } @@ -218,6 +255,10 @@ impl IggyBenchArgs { self.warmup_time } + pub const fn poll_wait_timeout(&self) -> IggyDuration { + self.poll_wait_timeout + } + pub const fn sampling_time(&self) -> IggyDuration { self.sampling_time } @@ -389,6 +430,10 @@ impl IggyBenchArgs { transport.to_string(), ]; + if !self.poll_wait_timeout().is_zero() { + parts.push(format!("poll_wait_{}", self.poll_wait_timeout())); + } + if let Some(remark) = &self.remark() { parts.push(remark.clone()); } @@ -447,3 +492,30 @@ impl IggyBenchArgs { name } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn poll_wait_timeout_defaults_to_zero() { + let args = IggyBenchArgs::try_parse_from(["iggy-bench", "pinned-consumer", "tcp"]) + .expect("args should parse"); + + assert!(args.poll_wait_timeout().is_zero()); + } + + #[test] + fn poll_wait_timeout_arg_parses() { + let args = IggyBenchArgs::try_parse_from([ + "iggy-bench", + "--poll-wait-timeout", + "10ms", + "pinned-consumer", + "tcp", + ]) + .expect("args should parse"); + + assert_eq!(args.poll_wait_timeout().as_micros(), 10_000); + } +} diff --git a/core/bench/src/args/defaults.rs b/core/bench/src/args/defaults.rs index a5f4dfac89..2aa127016d 100644 --- a/core/bench/src/args/defaults.rs +++ b/core/bench/src/args/defaults.rs @@ -46,6 +46,7 @@ pub const DEFAULT_NUMBER_OF_CONSUMER_GROUPS: NonZeroU32 = u32!(1); pub const DEFAULT_NUMBER_OF_PRODUCERS: NonZeroU32 = u32!(8); pub const DEFAULT_WARMUP_TIME: &str = "0s"; +pub const DEFAULT_POLL_WAIT_TIMEOUT: &str = "0s"; pub const DEFAULT_SAMPLING_TIME: &str = "10ms"; pub const DEFAULT_MOVING_AVERAGE_WINDOW: u32 = 20; diff --git a/core/bench/src/benchmarks/benchmark.rs b/core/bench/src/benchmarks/benchmark.rs index d1724e5a1f..136a86beae 100644 --- a/core/bench/src/benchmarks/benchmark.rs +++ b/core/bench/src/benchmarks/benchmark.rs @@ -193,7 +193,12 @@ pub trait Benchmarkable: Send { .rate_limit() .map(|rl| format!(" global rate limit: {rl}/s")) .unwrap_or_default(); + let poll_wait_timeout = if self.args().poll_wait_timeout().is_zero() { + String::new() + } else { + format!(" poll wait timeout: {},", self.args().poll_wait_timeout()) + }; - format!("{message_size}{messages_per_batch}{data}{rate_limit}") + format!("{message_size}{messages_per_batch}{data}{poll_wait_timeout}{rate_limit}") } } diff --git a/core/bench/src/benchmarks/common.rs b/core/bench/src/benchmarks/common.rs index cb5ee9f9a2..4f8ec0e602 100644 --- a/core/bench/src/benchmarks/common.rs +++ b/core/bench/src/benchmarks/common.rs @@ -165,6 +165,7 @@ pub fn build_consumer_futures( let consumers = args.consumers(); let actors = args.producers() + args.consumers(); let warmup_time = args.warmup_time(); + let poll_wait_timeout = args.poll_wait_timeout(); let messages_per_batch = args.messages_per_batch(); let sampling_time = args.sampling_time(); let moving_average_window = args.moving_average_window(); @@ -233,6 +234,7 @@ pub fn build_consumer_futures( messages_per_batch, finish_condition, warmup_time, + poll_wait_timeout, sampling_time, moving_average_window, polling_kind, @@ -255,6 +257,7 @@ pub fn build_producing_consumers_futures( let streams = args.streams(); let partitions = args.number_of_partitions(); let warmup_time = args.warmup_time(); + let poll_wait_timeout = args.poll_wait_timeout(); let messages_per_batch = args.messages_per_batch(); let message_size = args.message_size(); let polling_kind = PollingKind::Offset; @@ -296,6 +299,7 @@ pub fn build_producing_consumers_futures( send_finish_condition.clone(), poll_finish_condition.clone(), warmup_time, + poll_wait_timeout, args_clone.sampling_time(), args_clone.moving_average_window(), rate_limit, @@ -320,6 +324,7 @@ pub fn build_producing_consumer_groups_futures( let partitions = args.number_of_partitions(); let cg_count = args.number_of_consumer_groups(); let warmup_time = args.warmup_time(); + let poll_wait_timeout = args.poll_wait_timeout(); let messages_per_batch = args.messages_per_batch(); let message_size = args.message_size(); let start_consumer_group_id = CONSUMER_GROUP_BASE_ID; @@ -399,6 +404,7 @@ pub fn build_producing_consumer_groups_futures( send_finish_condition, poll_finish_condition, warmup_time, + poll_wait_timeout, args_clone.sampling_time(), args_clone.moving_average_window(), rate_limit, diff --git a/core/bench/src/utils/mod.rs b/core/bench/src/utils/mod.rs index d4d9f37532..ad20a7973b 100644 --- a/core/bench/src/utils/mod.rs +++ b/core/bench/src/utils/mod.rs @@ -21,7 +21,7 @@ use bench_report::{ transport::BenchmarkTransport, }; use iggy::prelude::*; -use std::{fs, path::Path}; +use std::{fmt::Write, fs, path::Path}; use tracing::{error, info}; use crate::args::{ @@ -31,8 +31,8 @@ use crate::args::{ DEFAULT_HTTP_SERVER_ADDRESS, DEFAULT_MESSAGE_BATCHES, DEFAULT_MESSAGE_SIZE, DEFAULT_MESSAGES_PER_BATCH, DEFAULT_NUMBER_OF_CONSUMER_GROUPS, DEFAULT_NUMBER_OF_CONSUMERS, DEFAULT_NUMBER_OF_PRODUCERS, DEFAULT_PINNED_NUMBER_OF_PARTITIONS, - DEFAULT_PINNED_NUMBER_OF_STREAMS, DEFAULT_QUIC_SERVER_ADDRESS, DEFAULT_TCP_SERVER_ADDRESS, - DEFAULT_TOTAL_MESSAGES_SIZE, DEFAULT_WARMUP_TIME, + DEFAULT_PINNED_NUMBER_OF_STREAMS, DEFAULT_POLL_WAIT_TIMEOUT, DEFAULT_QUIC_SERVER_ADDRESS, + DEFAULT_TCP_SERVER_ADDRESS, DEFAULT_TOTAL_MESSAGES_SIZE, DEFAULT_WARMUP_TIME, }, }; @@ -140,7 +140,10 @@ pub fn params_from_args_and_metrics( consumer_groups.to_string(), ]; - let params_identifier = params_identifier.join("_"); + let mut params_identifier = params_identifier.join("_"); + if !args.poll_wait_timeout().is_zero() { + let _ = write!(params_identifier, "_poll_wait_{}", args.poll_wait_timeout()); + } BenchmarkParams { benchmark_kind, @@ -227,6 +230,13 @@ fn add_basic_arguments(parts: &mut Vec, args: &IggyBenchArgs) { if args.warmup_time().to_string() != DEFAULT_WARMUP_TIME { parts.push(format!("--warmup-time \'{}\'", args.warmup_time())); } + + if args.poll_wait_timeout().to_string() != DEFAULT_POLL_WAIT_TIMEOUT { + parts.push(format!( + "--poll-wait-timeout \'{}\'", + args.poll_wait_timeout() + )); + } } fn add_benchmark_kind_arguments(parts: &mut Vec, args: &IggyBenchArgs) { diff --git a/core/common/src/traits/binary_impls/messages.rs b/core/common/src/traits/binary_impls/messages.rs index 332a18f0f1..6b022d1aca 100644 --- a/core/common/src/traits/binary_impls/messages.rs +++ b/core/common/src/traits/binary_impls/messages.rs @@ -189,6 +189,7 @@ async fn resolve_partitioning( /// (round-robin) and send an explicit-partition poll. A coordinator fence /// rejection (stale assignment after a rebalance) triggers one re-sync + retry. #[cfg(feature = "vsr")] +#[allow(clippy::too_many_arguments)] async fn poll_group_messages( client: &B, stream_id: &Identifier, diff --git a/core/integration/tests/server/cg.rs b/core/integration/tests/server/cg.rs index 697b3ddd92..b4138f7ae0 100644 --- a/core/integration/tests/server/cg.rs +++ b/core/integration/tests/server/cg.rs @@ -21,6 +21,7 @@ use crate::server::scenarios::{ consumer_group_new_messages_after_restart_scenario, consumer_group_offset_cleanup_scenario, consumer_group_with_multiple_clients_polling_messages_scenario, consumer_group_with_single_client_polling_messages_scenario, + poll_messages_wait_timeout_scenario, }; use integration::iggy_harness; @@ -81,3 +82,12 @@ async fn offset_cleanup(harness: &TestHarness) { async fn duplicate_name_create_preserves_live_group(harness: &TestHarness) { consumer_group_duplicate_name_create_scenario::run(harness).await; } + +#[cfg(not(feature = "vsr"))] +#[iggy_harness( + test_client_transport = [Tcp, WebSocket, Quic], + server(tcp.socket.override_defaults = true, tcp.socket.nodelay = true) +)] +async fn poll_messages_wait_timeout_consumer_group(harness: &TestHarness) { + poll_messages_wait_timeout_scenario::run_consumer_group_checks(harness).await; +} diff --git a/core/integration/tests/server/scenarios/mod.rs b/core/integration/tests/server/scenarios/mod.rs index 5d701d15e5..d5e7e75500 100644 --- a/core/integration/tests/server/scenarios/mod.rs +++ b/core/integration/tests/server/scenarios/mod.rs @@ -42,6 +42,7 @@ pub mod message_headers_scenario; pub mod message_size_scenario; pub mod offset_scenario; pub mod permissions_scenario; +pub mod poll_messages_wait_timeout_scenario; pub mod purge_delete_scenario; pub mod read_during_persistence_scenario; pub mod reconnect_after_restart_scenario; diff --git a/core/integration/tests/server/scenarios/poll_messages_wait_timeout_scenario.rs b/core/integration/tests/server/scenarios/poll_messages_wait_timeout_scenario.rs new file mode 100644 index 0000000000..204ca78c2a --- /dev/null +++ b/core/integration/tests/server/scenarios/poll_messages_wait_timeout_scenario.rs @@ -0,0 +1,429 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use iggy::prelude::*; +use integration::harness::TestHarness; +use std::time::{Duration, Instant}; +use tokio::time::sleep; +#[cfg(not(feature = "vsr"))] +use tokio::time::timeout; + +const STREAM_NAME: &str = "poll-wait-timeout-stream"; +const TOPIC_NAME: &str = "poll-wait-timeout-topic"; +const CONSUMER_GROUP_NAME: &str = "poll-wait-timeout-group"; +#[cfg(not(feature = "vsr"))] +const PARTITION_ID: u32 = 0; + +#[cfg(not(feature = "vsr"))] +pub async fn run_semantics_checks(harness: &TestHarness) { + let client = harness.root_client().await.expect("root client"); + let producer = harness.root_client().await.expect("producer client"); + setup_topic(&client, 1).await; + + verify_immediate_empty_polling(&client).await; + verify_timeout_waits_and_returns_empty(&client).await; + verify_existing_messages_return_without_waiting(&producer, &client).await; + verify_timeout_auto_commit_boundaries(&producer, &client).await; + + client + .delete_stream(&Identifier::named(STREAM_NAME).unwrap()) + .await + .unwrap(); +} + +#[cfg(not(feature = "vsr"))] +pub async fn run_wake_after_append_checks(harness: &TestHarness) { + let producer = harness.root_client().await.expect("producer client"); + let consumer_client = harness.root_client().await.expect("consumer client"); + + producer.create_stream(STREAM_NAME).await.unwrap(); + producer + .create_topic( + &Identifier::named(STREAM_NAME).unwrap(), + TOPIC_NAME, + 1, + CompressionAlgorithm::default(), + None, + IggyExpiry::NeverExpire, + MaxTopicSize::ServerDefault, + ) + .await + .unwrap(); + + let timeout = Duration::from_secs(5); + let start = Instant::now(); + let stream = Identifier::named(STREAM_NAME).unwrap(); + let topic = Identifier::named(TOPIC_NAME).unwrap(); + let consumer = Consumer::default(); + let strategy = PollingStrategy::offset(0); + let poll = consumer_client.poll_messages_with_timeout( + &stream, + &topic, + Some(0), + &consumer, + &strategy, + 1, + false, + timeout, + ); + let send = async { + sleep(Duration::from_millis(100)).await; + let mut messages = vec![ + IggyMessage::builder() + .id(1) + .payload("wake-after-append".into()) + .build() + .unwrap(), + ]; + producer + .send_messages( + &Identifier::named(STREAM_NAME).unwrap(), + &Identifier::named(TOPIC_NAME).unwrap(), + &Partitioning::partition_id(PARTITION_ID), + &mut messages, + ) + .await + .unwrap(); + }; + + let (polled, _) = tokio::join!(poll, send); + let polled = polled.unwrap(); + + assert_eq!(polled.partition_id, PARTITION_ID); + assert_eq!(polled.messages.len(), 1); + assert!( + start.elapsed() < timeout / 2, + "poll should wake after a successful append instead of waiting for timeout" + ); + + producer + .delete_stream(&Identifier::named(STREAM_NAME).unwrap()) + .await + .unwrap(); +} + +#[cfg(not(feature = "vsr"))] +async fn setup_topic(client: &IggyClient, partitions_count: u32) { + client.create_stream(STREAM_NAME).await.unwrap(); + client + .create_topic( + &Identifier::named(STREAM_NAME).unwrap(), + TOPIC_NAME, + partitions_count, + CompressionAlgorithm::default(), + None, + IggyExpiry::NeverExpire, + MaxTopicSize::ServerDefault, + ) + .await + .unwrap(); +} + +#[cfg(not(feature = "vsr"))] +async fn verify_immediate_empty_polling(client: &IggyClient) { + let stream = Identifier::named(STREAM_NAME).unwrap(); + let topic = Identifier::named(TOPIC_NAME).unwrap(); + let consumer = Consumer::new(Identifier::numeric(101).unwrap()); + let strategy = PollingStrategy::offset(0); + + let immediate = timeout( + Duration::from_secs(1), + client.poll_messages( + &stream, + &topic, + Some(PARTITION_ID), + &consumer, + &strategy, + 1, + false, + ), + ) + .await + .expect("immediate poll should not wait") + .unwrap(); + assert!(immediate.messages.is_empty()); + + let zero_timeout = timeout( + Duration::from_secs(1), + client.poll_messages_with_timeout( + &stream, + &topic, + Some(PARTITION_ID), + &consumer, + &strategy, + 1, + false, + Duration::ZERO, + ), + ) + .await + .expect("zero timeout poll should not wait") + .unwrap(); + assert!(zero_timeout.messages.is_empty()); +} + +#[cfg(not(feature = "vsr"))] +async fn verify_timeout_waits_and_returns_empty(client: &IggyClient) { + let stream = Identifier::named(STREAM_NAME).unwrap(); + let topic = Identifier::named(TOPIC_NAME).unwrap(); + let consumer = Consumer::new(Identifier::numeric(102).unwrap()); + let strategy = PollingStrategy::offset(0); + let poll = client.poll_messages_with_timeout( + &stream, + &topic, + Some(PARTITION_ID), + &consumer, + &strategy, + 1, + false, + Duration::from_millis(120), + ); + tokio::pin!(poll); + + tokio::select! { + result = &mut poll => panic!("non-zero timeout returned before its wait window: {result:?}"), + () = sleep(Duration::from_millis(20)) => {} + } + + let polled = poll.await.unwrap(); + assert!(polled.messages.is_empty()); +} + +#[cfg(not(feature = "vsr"))] +async fn verify_existing_messages_return_without_waiting( + producer: &IggyClient, + client: &IggyClient, +) { + let stream = Identifier::named(STREAM_NAME).unwrap(); + let topic = Identifier::named(TOPIC_NAME).unwrap(); + send_one(producer, 1, "already-readable").await; + + let polled = timeout( + Duration::from_secs(1), + client.poll_messages_with_timeout( + &stream, + &topic, + Some(PARTITION_ID), + &Consumer::new(Identifier::numeric(103).unwrap()), + &PollingStrategy::offset(0), + 1, + false, + Duration::from_secs(5), + ), + ) + .await + .expect("readable data should return without waiting for timeout") + .unwrap(); + + assert_eq!(polled.partition_id, PARTITION_ID); + assert_eq!(polled.messages.len(), 1); + assert_eq!(polled.messages[0].payload.as_ref(), b"already-readable"); +} + +#[cfg(not(feature = "vsr"))] +async fn verify_timeout_auto_commit_boundaries(producer: &IggyClient, client: &IggyClient) { + let stream = Identifier::named(STREAM_NAME).unwrap(); + let topic = Identifier::named(TOPIC_NAME).unwrap(); + let empty_timeout_consumer = Consumer::new(Identifier::numeric(104).unwrap()); + let empty = client + .poll_messages_with_timeout( + &stream, + &topic, + Some(PARTITION_ID), + &empty_timeout_consumer, + &PollingStrategy::offset(1), + 1, + true, + Duration::from_millis(50), + ) + .await + .unwrap(); + + assert!(empty.messages.is_empty()); + let offset_after_empty = client + .get_consumer_offset(&empty_timeout_consumer, &stream, &topic, Some(PARTITION_ID)) + .await + .unwrap(); + assert!(offset_after_empty.is_none()); + + let auto_commit_consumer = Consumer::new(Identifier::numeric(105).unwrap()); + let auto_commit_strategy = PollingStrategy::offset(1); + let poll = client.poll_messages_with_timeout( + &stream, + &topic, + Some(PARTITION_ID), + &auto_commit_consumer, + &auto_commit_strategy, + 1, + true, + Duration::from_secs(5), + ); + let send = async { + sleep(Duration::from_millis(100)).await; + send_one(producer, 2, "wake-auto-commit").await; + }; + let (polled, _) = tokio::join!(poll, send); + let polled = polled.unwrap(); + + assert_eq!(polled.messages.len(), 1); + let committed = client + .get_consumer_offset(&auto_commit_consumer, &stream, &topic, Some(PARTITION_ID)) + .await + .unwrap() + .expect("wake with returned messages should auto-commit"); + assert_eq!(committed.stored_offset, polled.messages[0].header.offset); +} + +#[cfg(not(feature = "vsr"))] +async fn send_one(client: &IggyClient, id: u128, payload: &str) { + let mut messages = vec![ + IggyMessage::builder() + .id(id) + .payload(payload.to_owned().into()) + .build() + .unwrap(), + ]; + client + .send_messages( + &Identifier::named(STREAM_NAME).unwrap(), + &Identifier::named(TOPIC_NAME).unwrap(), + &Partitioning::partition_id(PARTITION_ID), + &mut messages, + ) + .await + .unwrap(); +} + +pub async fn run_consumer_group_checks(harness: &TestHarness) { + let client = harness.root_client().await.expect("root client"); + + client.create_stream(STREAM_NAME).await.unwrap(); + client + .create_topic( + &Identifier::named(STREAM_NAME).unwrap(), + TOPIC_NAME, + 2, + CompressionAlgorithm::default(), + None, + IggyExpiry::NeverExpire, + MaxTopicSize::ServerDefault, + ) + .await + .unwrap(); + client + .create_consumer_group( + &Identifier::named(STREAM_NAME).unwrap(), + &Identifier::named(TOPIC_NAME).unwrap(), + CONSUMER_GROUP_NAME, + ) + .await + .unwrap(); + client + .join_consumer_group( + &Identifier::named(STREAM_NAME).unwrap(), + &Identifier::named(TOPIC_NAME).unwrap(), + &Identifier::named(CONSUMER_GROUP_NAME).unwrap(), + ) + .await + .unwrap(); + + let mut messages = vec![ + IggyMessage::builder() + .id(1) + .payload("ready-on-partition-1".into()) + .build() + .unwrap(), + ]; + client + .send_messages( + &Identifier::named(STREAM_NAME).unwrap(), + &Identifier::named(TOPIC_NAME).unwrap(), + &Partitioning::partition_id(1), + &mut messages, + ) + .await + .unwrap(); + + let consumer = Consumer::group(Identifier::named(CONSUMER_GROUP_NAME).unwrap()); + let group = client + .get_consumer_group( + &Identifier::named(STREAM_NAME).unwrap(), + &Identifier::named(TOPIC_NAME).unwrap(), + &Identifier::named(CONSUMER_GROUP_NAME).unwrap(), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(group.members[0].partitions, vec![0, 1]); + + wait_until_partition_readable(&client).await; + + let timeout = Duration::from_millis(750); + let start = Instant::now(); + let polled = client + .poll_messages_with_timeout( + &Identifier::named(STREAM_NAME).unwrap(), + &Identifier::named(TOPIC_NAME).unwrap(), + Some(0), + &consumer, + &PollingStrategy::offset(0), + 1, + true, + timeout, + ) + .await + .unwrap(); + + assert_eq!(polled.partition_id, 1); + assert_eq!(polled.messages.len(), 1); + assert!( + start.elapsed() < timeout / 2, + "poll should not wait on an empty owned partition while another owned partition has data" + ); + + client + .delete_stream(&Identifier::named(STREAM_NAME).unwrap()) + .await + .unwrap(); +} + +async fn wait_until_partition_readable(client: &IggyClient) { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + let direct = client + .poll_messages( + &Identifier::named(STREAM_NAME).unwrap(), + &Identifier::named(TOPIC_NAME).unwrap(), + Some(1), + &Consumer::default(), + &PollingStrategy::offset(0), + 1, + false, + ) + .await + .unwrap(); + + if direct.messages.len() == 1 { + return; + } + assert!( + Instant::now() < deadline, + "partition 1 should become readable" + ); + sleep(Duration::from_millis(25)).await; + } +} diff --git a/core/integration/tests/server/specific.rs b/core/integration/tests/server/specific.rs index 0c455d045e..d64ac4c096 100644 --- a/core/integration/tests/server/specific.rs +++ b/core/integration/tests/server/specific.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +#[cfg(not(feature = "vsr"))] +use crate::server::scenarios::poll_messages_wait_timeout_scenario; use crate::server::scenarios::{message_size_scenario, single_message_per_batch_scenario}; use crate::server::scenarios::{reconnect_after_restart_scenario, restart_offset_skip_scenario}; use crate::server::scenarios::{ @@ -59,6 +61,24 @@ async fn should_handle_single_message_per_batch_with_delayed_persistence(harness single_message_per_batch_scenario::run(harness, 5).await; } +#[iggy_harness( + test_client_transport = [Tcp, WebSocket, Quic], + server(tcp.socket.override_defaults = true, tcp.socket.nodelay = true) +)] +async fn poll_messages_wait_timeout_wakes_after_append(harness: &TestHarness) { + poll_messages_wait_timeout_scenario::run_wake_after_append_checks(harness).await; +} + +#[cfg(not(feature = "vsr"))] +#[iggy_harness( + test_client_transport = [Tcp], + server(tcp.socket.override_defaults = true, tcp.socket.nodelay = true) +)] +async fn poll_messages_wait_timeout_semantics(harness: &TestHarness) { + poll_messages_wait_timeout_scenario::run_semantics_checks(harness).await; +} + +#[cfg(not(feature = "vsr"))] #[iggy_harness( test_client_transport = [Tcp, WebSocket, Quic], server( diff --git a/core/server/src/binary/handlers/messages/poll_messages_handler.rs b/core/server/src/binary/handlers/messages/poll_messages_handler.rs index 2ac28ecfed..de22934e54 100644 --- a/core/server/src/binary/handlers/messages/poll_messages_handler.rs +++ b/core/server/src/binary/handlers/messages/poll_messages_handler.rs @@ -21,9 +21,13 @@ use crate::binary::dispatch::{ use crate::sender::SenderKind; use crate::shard::IggyShard; use crate::shard::system::messages::PollingArgs; +use crate::shard::transmission::message::ResolvedTopic; +use crate::shard::waiters::PollWaiterRegistration; +use crate::streaming::segments::IggyMessagesBatchSet; use crate::streaming::session::Session; +use futures::future::select_all; use iggy_binary_protocol::requests::messages::PollMessagesRequest; -use iggy_common::IggyError; +use iggy_common::{Consumer, IggyError, IggyPollMetadata, PollingStrategy}; use server_common::PooledBuffer; use server_common::sharding::IggyNamespace; use std::{rc::Rc, time::Duration}; @@ -58,41 +62,56 @@ pub async fn handle_poll_messages( topic, consumer.clone(), partition_id, - PollingArgs::with_wait_timeout(strategy.clone(), count, auto_commit, wait_timeout), + PollingArgs::with_wait_timeout(strategy, count, auto_commit, wait_timeout), ) .await?; - if !wait_timeout.is_zero() && batch.is_empty() { - let namespace = IggyNamespace::new( - topic.stream_id, - topic.topic_id, - metadata.partition_id as usize, - ); + if count > 0 && !wait_timeout.is_zero() && batch.is_empty() { + let namespaces = + shard.resolve_poll_wait_namespaces(topic, &consumer, client_id, partition_id)?; + let waiters = namespaces + .iter() + .filter_map(|namespace| shard.register_poll_waiter(*namespace, wait_timeout)) + .collect::>(); - if let Some(waiter) = shard.register_poll_waiter(namespace, wait_timeout) { - let immediate_args = || PollingArgs::new(strategy.clone(), count, auto_commit); - (metadata, batch) = shard - .poll_messages( + if let Some((next_metadata, next_batch)) = poll_wait_namespaces( + shard, + client_id, + topic, + &consumer, + &namespaces, + &strategy, + count, + auto_commit, + ) + .await? + { + metadata = next_metadata; + batch = next_batch; + } + + if batch.is_empty() && !waiters.is_empty() { + let woke = matches!( + compio::time::timeout(wait_timeout, wait_for_any_poll_waiter(&waiters)).await, + Ok(true) + ); + drop(waiters); + + if woke + && let Some((next_metadata, next_batch)) = poll_wait_namespaces( + shard, client_id, topic, - consumer.clone(), - partition_id, - immediate_args(), + &consumer, + &namespaces, + &strategy, + count, + auto_commit, ) - .await?; - - if batch.is_empty() { - let woke = matches!( - compio::time::timeout(wait_timeout, waiter.wait()).await, - Ok(true) - ); - drop(waiter); - - if woke { - (metadata, batch) = shard - .poll_messages(client_id, topic, consumer, partition_id, immediate_args()) - .await?; - } + .await? + { + metadata = next_metadata; + batch = next_batch; } } } @@ -126,3 +145,39 @@ pub async fn handle_poll_messages( .await?; Ok(HandlerResult::Finished) } + +#[allow(clippy::too_many_arguments)] +async fn poll_wait_namespaces( + shard: &IggyShard, + client_id: u32, + topic: ResolvedTopic, + consumer: &Consumer, + namespaces: &[IggyNamespace], + strategy: &PollingStrategy, + count: u32, + auto_commit: bool, +) -> Result, IggyError> { + for namespace in namespaces { + let result = shard + .poll_messages( + client_id, + topic, + consumer.clone(), + Some(namespace.partition_id() as u32), + PollingArgs::new(*strategy, count, auto_commit), + ) + .await?; + if !result.1.is_empty() { + return Ok(Some(result)); + } + } + Ok(None) +} + +async fn wait_for_any_poll_waiter(waiters: &[PollWaiterRegistration]) -> bool { + let waits = waiters + .iter() + .map(|waiter| Box::pin(waiter.wait())) + .collect::>(); + select_all(waits).await.0 +} diff --git a/core/server/src/main.rs b/core/server/src/main.rs index 1f0ac18d90..2ae91c87f6 100644 --- a/core/server/src/main.rs +++ b/core/server/src/main.rs @@ -35,6 +35,7 @@ use server::log::logger::Logging; use server::metadata::{Metadata, create_metadata_handles}; use server::server_error::ServerError; use server::shard::system::info::SystemInfo; +use server::shard::waiters::PollWaiterRegistry; use server::shard::{IggyShard, calculate_shard_assignment}; use server::state::file::FileState; use server::state::system::SystemState; @@ -49,9 +50,9 @@ use shard_allocator::ShardAllocator; use std::panic::AssertUnwindSafe; use std::rc::Rc; use std::str::FromStr; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; use std::sync::mpsc; +use std::sync::{Arc, Mutex}; use std::thread::JoinHandle; use system_stats::capture_allowed_cpus; use tracing::{error, info, instrument, warn}; @@ -339,6 +340,7 @@ fn main() -> Result<(), ServerError> { let client_manager = Box::leak(client_manager); let client_manager: EternalPtr> = client_manager.into(); let client_manager = ClientManager::new(client_manager); + let poll_waiters = Arc::new(Mutex::new(PollWaiterRegistry::default())); // Populate shards_table from SharedMetadata partitions (hierarchical traversal) metadata.with_metadata(|metadata| { @@ -384,6 +386,7 @@ fn main() -> Result<(), ServerError> { state_term.clone(), ); let client_manager = client_manager.clone(); + let poll_waiters = poll_waiters.clone(); let shard_metadata = metadata.clone(); // Take metadata_writer for shard 0 only @@ -431,7 +434,8 @@ fn main() -> Result<(), ServerError> { .metrics(metrics) .is_follower(is_follower) .current_replica_id(replica_id) - .metadata(shard_metadata); + .metadata(shard_metadata) + .poll_waiters(poll_waiters); if let Some(writer) = shard_metadata_writer { builder = builder.metadata_writer(writer); diff --git a/core/server/src/shard/builder.rs b/core/server/src/shard/builder.rs index d8dc25006d..f7e41c3eb3 100644 --- a/core/server/src/shard/builder.rs +++ b/core/server/src/shard/builder.rs @@ -37,7 +37,7 @@ use server_common::sharding::{IggyNamespace, PartitionLocation}; use std::{ cell::{Cell, RefCell}, rc::Rc, - sync::atomic::AtomicBool, + sync::{Arc, Mutex, atomic::AtomicBool}, }; #[derive(Default)] @@ -58,6 +58,7 @@ pub struct IggyShardBuilder { current_replica_id: Option, metadata: Option, metadata_writer: Option, + poll_waiters: Option>>, } impl IggyShardBuilder { @@ -129,6 +130,11 @@ impl IggyShardBuilder { self } + pub fn poll_waiters(mut self, poll_waiters: Arc>) -> Self { + self.poll_waiters = Some(poll_waiters); + self + } + // TODO: Too much happens in there, some of those bootstrapping logic should be moved outside. pub fn build(self) -> IggyShard { let id = self.id.unwrap(); @@ -174,7 +180,9 @@ impl IggyShardBuilder { metadata_writer: self.metadata_writer.map(RefCell::new), local_partitions, pending_partition_inits: RefCell::new(AHashSet::new()), - poll_waiters: RefCell::new(PollWaiterRegistry::default()), + poll_waiters: self + .poll_waiters + .unwrap_or_else(|| Arc::new(Mutex::new(PollWaiterRegistry::default()))), encryptor, config, _version: version, diff --git a/core/server/src/shard/mod.rs b/core/server/src/shard/mod.rs index 66a35c4df1..a729f6b8bd 100644 --- a/core/server/src/shard/mod.rs +++ b/core/server/src/shard/mod.rs @@ -43,7 +43,7 @@ use std::{ net::SocketAddr, rc::Rc, sync::{ - Arc, + Arc, Mutex, atomic::{AtomicBool, AtomicU64, Ordering}, }, time::{Duration, Instant}, @@ -80,7 +80,7 @@ pub struct IggyShard { pub(crate) metadata_writer: Option>, pub(crate) local_partitions: RefCell, pub(crate) pending_partition_inits: RefCell>, - pub(crate) poll_waiters: RefCell, + pub(crate) poll_waiters: Arc>, pub(crate) shards_table: EternalPtr>, pub(crate) state: FileState, diff --git a/core/server/src/shard/system/utils.rs b/core/server/src/shard/system/utils.rs index fc5f37a5f7..594fee0056 100644 --- a/core/server/src/shard/system/utils.rs +++ b/core/server/src/shard/system/utils.rs @@ -26,6 +26,7 @@ use crate::{ streaming::polling_consumer::PollingConsumer, }; use iggy_common::{Consumer, ConsumerKind, Identifier, IggyError}; +use server_common::sharding::IggyNamespace; impl IggyShard { /// Resolves stream identifier to typed `ResolvedStream`. @@ -189,6 +190,60 @@ impl IggyShard { } } + pub fn resolve_poll_wait_namespaces( + &self, + topic: ResolvedTopic, + consumer: &Consumer, + client_id: u32, + partition_id: Option, + ) -> Result, IggyError> { + if consumer.kind == ConsumerKind::Consumer { + let Some((_, partition_id)) = self.resolve_consumer_with_partition_id( + topic, + consumer, + client_id, + partition_id, + false, + )? + else { + return Ok(Vec::new()); + }; + return Ok(vec![IggyNamespace::new( + topic.stream_id, + topic.topic_id, + partition_id, + )]); + } + + if self.client_manager.try_get_client(client_id).is_none() { + return Err(IggyError::StaleClient); + } + + let mut namespaces = Vec::new(); + for partition_id in self + .metadata + .get_partition_ids(topic.stream_id, topic.topic_id) + { + let Some((_, partition_id)) = self.metadata.resolve_consumer_group_partition( + topic.stream_id, + topic.topic_id, + &consumer.id, + client_id, + Some(partition_id as u32), + false, + )? + else { + continue; + }; + namespaces.push(IggyNamespace::new( + topic.stream_id, + topic.topic_id, + partition_id, + )); + } + Ok(namespaces) + } + /// Resolves topic and verifies user has append permission atomically. pub fn resolve_topic_for_append( &self, diff --git a/core/server/src/shard/waiters.rs b/core/server/src/shard/waiters.rs index 3ee68fd540..0af4e2b3e1 100644 --- a/core/server/src/shard/waiters.rs +++ b/core/server/src/shard/waiters.rs @@ -19,7 +19,7 @@ use crate::shard::IggyShard; use ahash::AHashMap; use async_channel::{Receiver, Sender}; use server_common::sharding::IggyNamespace; -use std::{cell::RefCell, time::Duration, time::Instant}; +use std::{sync::Arc, sync::Mutex, time::Duration, time::Instant}; const MAX_WAITERS_PER_NAMESPACE: usize = 1024; @@ -31,7 +31,7 @@ struct PollWaiter { } #[derive(Debug, Default)] -pub(crate) struct PollWaiterRegistry { +pub struct PollWaiterRegistry { next_id: u64, waiters: AHashMap>, } @@ -126,22 +126,25 @@ impl PollWaiterRegistry { } } -pub(crate) struct PollWaiterRegistration<'a> { +pub(crate) struct PollWaiterRegistration { namespace: IggyNamespace, id: u64, receiver: Receiver<()>, - registry: &'a RefCell, + registry: Arc>, } -impl PollWaiterRegistration<'_> { +impl PollWaiterRegistration { pub(crate) async fn wait(&self) -> bool { self.receiver.recv().await.is_ok() } } -impl Drop for PollWaiterRegistration<'_> { +impl Drop for PollWaiterRegistration { fn drop(&mut self) { - self.registry.borrow_mut().remove(&self.namespace, self.id); + self.registry + .lock() + .expect("poll waiter registry poisoned") + .remove(&self.namespace, self.id); } } @@ -150,30 +153,66 @@ impl IggyShard { &self, namespace: IggyNamespace, timeout: Duration, - ) -> Option> { + ) -> Option { let (id, receiver) = self .poll_waiters - .borrow_mut() + .lock() + .expect("poll waiter registry poisoned") .register(namespace, timeout)?; Some(PollWaiterRegistration { namespace, id, receiver, - registry: &self.poll_waiters, + registry: self.poll_waiters.clone(), }) } pub(crate) fn wake_poll_waiters(&self, namespace: &IggyNamespace) { - self.poll_waiters.borrow_mut().wake_namespace(namespace); + self.poll_waiters + .lock() + .expect("poll waiter registry poisoned") + .wake_namespace(namespace); } pub(crate) fn wake_topic_poll_waiters(&self, stream_id: usize, topic_id: usize) { self.poll_waiters - .borrow_mut() + .lock() + .expect("poll waiter registry poisoned") .wake_topic(stream_id, topic_id); } pub(crate) fn wake_stream_poll_waiters(&self, stream_id: usize) { - self.poll_waiters.borrow_mut().wake_stream(stream_id); + self.poll_waiters + .lock() + .expect("poll waiter registry poisoned") + .wake_stream(stream_id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn dropped_registration_removes_waiter() { + let registry = Arc::new(Mutex::new(PollWaiterRegistry::default())); + let namespace = IggyNamespace::new(1, 1, 0); + let (id, receiver) = registry + .lock() + .unwrap() + .register(namespace, Duration::from_secs(1)) + .expect("waiter should register"); + + assert_eq!(registry.lock().unwrap().waiters[&namespace].len(), 1); + + let registration = PollWaiterRegistration { + namespace, + id, + receiver, + registry: registry.clone(), + }; + drop(registration); + + assert!(!registry.lock().unwrap().waiters.contains_key(&namespace)); } } From 24026c5beef8dda230aa6ea6fe343168c7abb78a Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Sat, 4 Jul 2026 14:46:59 +0530 Subject: [PATCH 18/24] bench: add poll wait timeout comparison runner --- .../poll-wait-timeout-comparison.md | 66 +++++ .../run-poll-wait-timeout-comparison.sh | 233 ++++++++++++++++++ 2 files changed, 299 insertions(+) create mode 100644 scripts/performance/poll-wait-timeout-comparison.md create mode 100755 scripts/performance/run-poll-wait-timeout-comparison.sh diff --git a/scripts/performance/poll-wait-timeout-comparison.md b/scripts/performance/poll-wait-timeout-comparison.md new file mode 100644 index 0000000000..c46701d9db --- /dev/null +++ b/scripts/performance/poll-wait-timeout-comparison.md @@ -0,0 +1,66 @@ +# Poll Wait Timeout Benchmark Comparison + +This runbook compares immediate `PollMessages` polling with deferred polling. + +## Goal + +Compare immediate polling with deferred polling on sparse and saturated workloads. The script records `iggy-bench` throughput and latency; use OS-level sampling or future poll-count metrics for direct CPU/request-churn numbers. + +## What To Compare + +- Immediate polling: `--poll-wait-timeout 0s` +- Deferred polling: `--poll-wait-timeout 10ms`, `100ms`, `1s` +- Transports: start with `tcp`; optionally repeat with `websocket` +- Sparse workloads: low producer rate, small batches +- Saturated control: large batches, no rate limit + +## Run + +Quick local smoke, about 2-3 minutes on a warm release build: + +```bash +scripts/performance/run-poll-wait-timeout-comparison.sh --quick +``` + +Full local run: + +```bash +scripts/performance/run-poll-wait-timeout-comparison.sh +``` + +WebSocket-facing comparison: + +```bash +scripts/performance/run-poll-wait-timeout-comparison.sh --transport websocket +``` + +Print commands without running: + +```bash +scripts/performance/run-poll-wait-timeout-comparison.sh --dry-run +``` + +## Workloads + +The script runs: + +- `sparse_single`: one producer and one consumer, `1` message per batch, `10KB/s` rate limit +- `sparse_consumer_group`: one producer and four consumers over four partitions, `1` message per batch, `10KB/s` rate limit +- `saturated_control`: one producer and one consumer, `1000` messages per batch, no rate limit + +## Success Criteria + +- Deferred sparse runs should not lose messages. +- Sparse workloads should keep expected low-rate throughput while showing bounded wait-path latency. +- Direct CPU/request-churn claims should be backed by OS-level sampling or explicit poll-count metrics. +- Saturated control with non-zero timeout should be close to immediate polling for throughput and p99 latency. +- WebSocket comparison should behave consistently with TCP. + +## Discord Update Template + +```text +I added a focused benchmark comparison runner for #3605: +scripts/performance/run-poll-wait-timeout-comparison.sh + +It compares --poll-wait-timeout 0s vs 10ms/100ms/1s on sparse single-consumer, sparse consumer-group, and saturated control workloads. I’ll share the result table after the TCP quick/full run, and can repeat on WebSocket if useful. +``` diff --git a/scripts/performance/run-poll-wait-timeout-comparison.sh b/scripts/performance/run-poll-wait-timeout-comparison.sh new file mode 100755 index 0000000000..cf0c29eb0c --- /dev/null +++ b/scripts/performance/run-poll-wait-timeout-comparison.sh @@ -0,0 +1,233 @@ +#!/bin/bash +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# shellcheck disable=SC1091 + +set -euo pipefail + +IGGY_BENCH_CMD="target/release/iggy-bench" +IGGY_SERVER_CMD="target/release/iggy-server" +IDENTIFIER="$(hostname)" +OUTPUT_DIR="performance_results/poll_wait_timeout_comparison" +TRANSPORT="tcp" +SKIP_BUILD=false +DRY_RUN=false +QUICK=false + +while [[ $# -gt 0 ]]; do + case "$1" in + --bench-cmd) + IGGY_BENCH_CMD="$2" + shift 2 + ;; + --server-cmd) + IGGY_SERVER_CMD="$2" + shift 2 + ;; + --identifier) + IDENTIFIER="$2" + shift 2 + ;; + --output-dir) + OUTPUT_DIR="$2" + shift 2 + ;; + --transport) + TRANSPORT="$2" + shift 2 + ;; + --skip-build) + SKIP_BUILD=true + shift + ;; + --dry-run) + DRY_RUN=true + shift + ;; + --quick) + QUICK=true + shift + ;; + *) + echo "Unknown argument: $1" >&2 + exit 1 + ;; + esac +done + +case "$TRANSPORT" in +tcp | websocket | quic) + ;; +*) + echo "Unsupported transport: ${TRANSPORT}. Use tcp, websocket, or quic." >&2 + exit 1 + ;; +esac + +source "$(dirname "$0")/../utils.sh" +source "$(dirname "$0")/utils.sh" + +function on_interrupt() { + on_exit_bench + exit 130 +} + +if [[ "$DRY_RUN" == false ]]; then + trap on_interrupt SIGINT SIGTERM + trap on_exit_bench EXIT +fi + +if [[ "$SKIP_BUILD" == false && "$DRY_RUN" == false ]]; then + echo "Building release binaries..." + RUSTFLAGS="-C target-cpu=native" cargo build --release --bin iggy-server --bin iggy-bench +fi + +if [[ "$DRY_RUN" == false ]]; then + mkdir -p "$OUTPUT_DIR" +fi + +commit_hash=$(get_git_iggy_server_tag_or_sha1 .) +commit_date=$(get_git_commit_date .) + +if [[ "$QUICK" == true ]]; then + SPARSE_BATCHES=30 + SATURATED_BATCHES=20 + GROUP_BATCHES=20 +else + SPARSE_BATCHES=1000 + SATURATED_BATCHES=250 + GROUP_BATCHES=1000 +fi + +wait_timeouts=("0s" "10ms" "100ms" "1s") + +function bench_transport() { + case "$TRANSPORT" in + tcp) + echo "tcp" + ;; + websocket) + echo "websocket" + ;; + quic) + echo "quic" + ;; + esac +} + +function result_log_name() { + local suite="$1" + local wait_timeout="$2" + echo "${OUTPUT_DIR}/${suite}_wait_${wait_timeout//[^a-zA-Z0-9]/_}.log" +} + +function start_server() { + echo "Cleaning old local_data..." + rm -rf local_data + + echo "Starting iggy-server..." + IGGY_ROOT_USERNAME=iggy IGGY_ROOT_PASSWORD=iggy "$IGGY_SERVER_CMD" &>"${OUTPUT_DIR}/iggy-server.log" & + sleep 2 + + local server_pid + server_pid=$(pgrep -x "iggy-server") + exit_if_process_is_not_running "$server_pid" +} + +function stop_server() { + echo "Stopping iggy-server..." + send_signal "iggy-server" "TERM" + wait_for_process "iggy-server" 10 || send_signal "iggy-server" "KILL" +} + +function run_suite() { + local suite="$1" + local wait_timeout="$2" + local command="$3" + local log_file + log_file=$(result_log_name "$suite" "$wait_timeout") + + echo + echo "Suite: ${suite}, poll wait timeout: ${wait_timeout}" + echo "Command: ${command}" + + if [[ "$DRY_RUN" == true ]]; then + return 0 + fi + + start_server + set +e + eval "$command" 2>&1 | tee "$log_file" + local status=${PIPESTATUS[0]} + set -e + stop_server + + echo + echo "Summary from ${log_file}:" + grep -E "Results:|Total throughput|latency:" "$log_file" || true + + if [[ $status -ne 0 ]]; then + echo "Benchmark failed with status ${status}: ${suite}, wait=${wait_timeout}" >&2 + exit "$status" + fi +} + +function common_output_args() { + local suite="$1" + local wait_timeout="$2" + echo "output --output-dir ${OUTPUT_DIR} --identifier ${IDENTIFIER} --remark ${suite}_wait_${wait_timeout} --extra-info poll_wait_timeout_comparison --gitref ${commit_hash} --gitref-date ${commit_date}" +} + +function command_for_sparse_single() { + local wait_timeout="$1" + local output_args + output_args=$(common_output_args "sparse_single" "$wait_timeout") + echo "${IGGY_BENCH_CMD} --message-size 1000 --messages-per-batch 1 --message-batches ${SPARSE_BATCHES} --rate-limit 10KB --poll-wait-timeout ${wait_timeout} pinned-producer-and-consumer --streams 1 --producers 1 --consumers 1 $(bench_transport) ${output_args}" +} + +function command_for_sparse_consumer_group() { + local wait_timeout="$1" + local output_args + output_args=$(common_output_args "sparse_consumer_group" "$wait_timeout") + echo "${IGGY_BENCH_CMD} --message-size 1000 --messages-per-batch 1 --message-batches ${GROUP_BATCHES} --rate-limit 10KB --poll-wait-timeout ${wait_timeout} balanced-producer-and-consumer-group --streams 1 --partitions 4 --producers 1 --consumers 4 $(bench_transport) ${output_args}" +} + +function command_for_saturated_control() { + local wait_timeout="$1" + local output_args + output_args=$(common_output_args "saturated_control" "$wait_timeout") + echo "${IGGY_BENCH_CMD} --message-size 1000 --messages-per-batch 1000 --message-batches ${SATURATED_BATCHES} --poll-wait-timeout ${wait_timeout} pinned-producer-and-consumer --streams 1 --producers 1 --consumers 1 $(bench_transport) ${output_args}" +} + +echo "Running PollMessages wait-timeout benchmark comparison" +echo "transport=${TRANSPORT}, identifier=${IDENTIFIER}, output_dir=${OUTPUT_DIR}, quick=${QUICK}, dry_run=${DRY_RUN}" + +for wait_timeout in "${wait_timeouts[@]}"; do + run_suite "sparse_single" "$wait_timeout" "$(command_for_sparse_single "$wait_timeout")" +done + +for wait_timeout in "${wait_timeouts[@]}"; do + run_suite "sparse_consumer_group" "$wait_timeout" "$(command_for_sparse_consumer_group "$wait_timeout")" +done + +for wait_timeout in "0s" "100ms"; do + run_suite "saturated_control" "$wait_timeout" "$(command_for_saturated_control "$wait_timeout")" +done + +echo +echo "Comparison complete. Full logs and benchmark artifacts are in ${OUTPUT_DIR}." From 6213fea7f4ca1b70293064c7ad3accb54a9e2fcd Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Sun, 5 Jul 2026 17:41:42 +0530 Subject: [PATCH 19/24] chore(bench): extend poll timeout comparison --- .../poll-wait-timeout-comparison.md | 79 +++++++-- .../run-poll-wait-timeout-comparison.sh | 159 +++++++++++++++++- 2 files changed, 222 insertions(+), 16 deletions(-) diff --git a/scripts/performance/poll-wait-timeout-comparison.md b/scripts/performance/poll-wait-timeout-comparison.md index c46701d9db..72b7ffe0be 100644 --- a/scripts/performance/poll-wait-timeout-comparison.md +++ b/scripts/performance/poll-wait-timeout-comparison.md @@ -4,36 +4,61 @@ This runbook compares immediate `PollMessages` polling with deferred polling. ## Goal -Compare immediate polling with deferred polling on sparse and saturated workloads. The script records `iggy-bench` throughput and latency; use OS-level sampling or future poll-count metrics for direct CPU/request-churn numbers. +Compare immediate polling with deferred polling on sparse, busy-loop, consumer-group, and saturated workloads. The script records `iggy-bench` throughput and latency, plus OS-level CPU/RSS samples and network byte deltas for each run. ## What To Compare - Immediate polling: `--poll-wait-timeout 0s` - Deferred polling: `--poll-wait-timeout 10ms`, `100ms`, `1s` -- Transports: start with `tcp`; optionally repeat with `websocket` +- Busy-loop comparison: `0s` vs `100ms` +- Transports: run `tcp` first, then repeat with `websocket` - Sparse workloads: low producer rate, small batches - Saturated control: large batches, no rate limit +- Resource signals: `iggy-server` CPU/RSS, `iggy-bench` CPU/RSS, network bytes before/after each run ## Run -Quick local smoke, about 2-3 minutes on a warm release build: +Quick TCP smoke, about 2-3 minutes on a warm release build: ```bash -scripts/performance/run-poll-wait-timeout-comparison.sh --quick +scripts/performance/run-poll-wait-timeout-comparison.sh --quick --transport tcp ``` -Full local run: +Full TCP run: ```bash -scripts/performance/run-poll-wait-timeout-comparison.sh +scripts/performance/run-poll-wait-timeout-comparison.sh --transport tcp ``` -WebSocket-facing comparison: +Full WebSocket run: ```bash scripts/performance/run-poll-wait-timeout-comparison.sh --transport websocket ``` +Use existing release binaries when you already built them elsewhere: + +```bash +scripts/performance/run-poll-wait-timeout-comparison.sh \ + --quick \ + --skip-build \ + --bench-cmd /Users/aruns/Developer/iggy/target/release/iggy-bench \ + --server-cmd /Users/aruns/Developer/iggy/target/release/iggy-server \ + --identifier aruns-mbp-pr3605 +``` + +Change OS sample interval: + +```bash +scripts/performance/run-poll-wait-timeout-comparison.sh --sample-interval 0.5 +``` + +Disable OS sampling if you only want `iggy-bench` throughput and latency: + +```bash +scripts/performance/run-poll-wait-timeout-comparison.sh --no-resource-sampling +``` + Print commands without running: ```bash @@ -45,22 +70,50 @@ scripts/performance/run-poll-wait-timeout-comparison.sh --dry-run The script runs: - `sparse_single`: one producer and one consumer, `1` message per batch, `10KB/s` rate limit +- `busy_loop_single`: one producer and one consumer, `1` message per batch, `5KB/s` rate limit, `0s` vs `100ms` - `sparse_consumer_group`: one producer and four consumers over four partitions, `1` message per batch, `10KB/s` rate limit -- `saturated_control`: one producer and one consumer, `1000` messages per batch, no rate limit +- `saturated_control`: one producer and one consumer, `1000` messages per batch, no rate limit, `0s` vs `100ms` + +## Resource Logs + +Each run writes a benchmark log and, when resource sampling is enabled, a matching resource CSV: + +```text +performance_results/poll_wait_timeout_comparison/_wait_.log +performance_results/poll_wait_timeout_comparison/_wait__resources.csv +``` + +The CSV format is: + +```text +timestamp,role,pid,cpu_percent,rss_kb +``` + +The benchmark log appends: + +- average and max CPU for `server` and `bench` +- max RSS for `server` and `bench` +- network byte delta for the run + +## Interpretation Caveats + +- Local TCP and WebSocket runs usually travel over loopback, so network bytes are useful for relative comparison, not external NIC throughput. +- OS CPU sampling is interval-based and process-level; use it to compare `0s` vs deferred polling trends, not as a microbenchmark-grade profiler. +- Busy-loop value should show up most clearly in CPU and network/request churn, especially when producer rate is sparse. +- Saturated control protects against regression when messages are already readable; deferred polling should stay close to immediate polling for throughput and p99 latency. +- Repeat quick smoke before sharing numbers; use full TCP/WebSocket runs for final PR discussion data. ## Success Criteria - Deferred sparse runs should not lose messages. - Sparse workloads should keep expected low-rate throughput while showing bounded wait-path latency. -- Direct CPU/request-churn claims should be backed by OS-level sampling or explicit poll-count metrics. +- Busy-loop deferred runs should reduce CPU/network churn relative to `0s` immediate polling. - Saturated control with non-zero timeout should be close to immediate polling for throughput and p99 latency. - WebSocket comparison should behave consistently with TCP. +- Resource claims should cite the generated CSV/log data and avoid overclaiming precision. ## Discord Update Template ```text -I added a focused benchmark comparison runner for #3605: -scripts/performance/run-poll-wait-timeout-comparison.sh - -It compares --poll-wait-timeout 0s vs 10ms/100ms/1s on sparse single-consumer, sparse consumer-group, and saturated control workloads. I’ll share the result table after the TCP quick/full run, and can repeat on WebSocket if useful. +I extended the #3605 benchmark runner to cover regular polling vs deferred polling in busy-loop/sparse workloads, plus OS-level CPU/RSS samples and network byte deltas per run. I’ll run TCP first, then repeat WebSocket, and share the 0s vs 100ms comparison with throughput, p99, CPU, RSS, and network deltas. ``` diff --git a/scripts/performance/run-poll-wait-timeout-comparison.sh b/scripts/performance/run-poll-wait-timeout-comparison.sh index cf0c29eb0c..a322780ecc 100755 --- a/scripts/performance/run-poll-wait-timeout-comparison.sh +++ b/scripts/performance/run-poll-wait-timeout-comparison.sh @@ -28,6 +28,8 @@ TRANSPORT="tcp" SKIP_BUILD=false DRY_RUN=false QUICK=false +RESOURCE_SAMPLING=true +SAMPLE_INTERVAL_SECONDS=1 while [[ $# -gt 0 ]]; do case "$1" in @@ -51,6 +53,14 @@ while [[ $# -gt 0 ]]; do TRANSPORT="$2" shift 2 ;; + --sample-interval) + SAMPLE_INTERVAL_SECONDS="$2" + shift 2 + ;; + --no-resource-sampling) + RESOURCE_SAMPLING=false + shift + ;; --skip-build) SKIP_BUILD=true shift @@ -79,6 +89,11 @@ tcp | websocket | quic) ;; esac +if ! [[ "$SAMPLE_INTERVAL_SECONDS" =~ ^[0-9]+([.][0-9]+)?$ ]]; then + echo "Invalid sample interval: ${SAMPLE_INTERVAL_SECONDS}" >&2 + exit 1 +fi + source "$(dirname "$0")/../utils.sh" source "$(dirname "$0")/utils.sh" @@ -108,13 +123,16 @@ if [[ "$QUICK" == true ]]; then SPARSE_BATCHES=30 SATURATED_BATCHES=20 GROUP_BATCHES=20 + BUSY_LOOP_BATCHES=20 else SPARSE_BATCHES=1000 SATURATED_BATCHES=250 GROUP_BATCHES=1000 + BUSY_LOOP_BATCHES=200 fi wait_timeouts=("0s" "10ms" "100ms" "1s") +busy_loop_wait_timeouts=("0s" "100ms") function bench_transport() { case "$TRANSPORT" in @@ -136,6 +154,108 @@ function result_log_name() { echo "${OUTPUT_DIR}/${suite}_wait_${wait_timeout//[^a-zA-Z0-9]/_}.log" } +function snapshot_network_bytes() { + local platform + platform=$(uname -s) + + if [[ "$platform" == "Darwin" ]]; then + netstat -ibn 2>/dev/null | awk 'NR > 1 { rx += $7; tx += $10 } END { printf "%.0f %.0f\n", rx, tx }' + return + fi + + if [[ -d /sys/class/net ]]; then + local rx=0 + local tx=0 + local iface + for iface in /sys/class/net/*; do + if [[ -r "$iface/statistics/rx_bytes" && -r "$iface/statistics/tx_bytes" ]]; then + rx=$((rx + $(<"$iface/statistics/rx_bytes"))) + tx=$((tx + $(<"$iface/statistics/tx_bytes"))) + fi + done + echo "$rx $tx" + return + fi + + echo "0 0" +} + +function sample_role_resources() { + local role="$1" + local process_name="$2" + local resource_log="$3" + local pids + local pid + local line + + pids=$(pgrep -x "$process_name" || true) + for pid in $pids; do + line=$(ps -p "$pid" -o pid= -o %cpu= -o rss= 2>/dev/null | awk '{$1=$1; print}' || true) + if [[ -n "$line" ]]; then + awk -v ts="$(date +%s)" -v role="$role" '{ print ts "," role "," $1 "," $2 "," $3 }' <<<"$line" >>"$resource_log" + fi + done +} + +function sample_process_resources() { + local resource_log="$1" + + while true; do + sample_role_resources "server" "iggy-server" "$resource_log" + sample_role_resources "bench" "iggy-bench" "$resource_log" + sleep "$SAMPLE_INTERVAL_SECONDS" + done +} + +function stop_resource_sampler() { + local sampler_pid="$1" + + if [[ -n "$sampler_pid" ]]; then + kill "$sampler_pid" 2>/dev/null || true + wait "$sampler_pid" 2>/dev/null || true + fi +} + +function summarize_resource_role() { + local resource_log="$1" + local role="$2" + + awk -F, -v role="$role" ' + $2 == role { + count++; + cpu += $4; + if ($4 > max_cpu) { max_cpu = $4; } + if ($5 > max_rss) { max_rss = $5; } + } + END { + if (count > 0) { + printf "Resource summary: %s avg_cpu=%.2f%% max_cpu=%.2f%% max_rss_kb=%d samples=%d\n", role, cpu / count, max_cpu, max_rss, count; + } else { + printf "Resource summary: %s no samples captured\n", role; + } + } + ' "$resource_log" +} + +function append_resource_summary() { + local log_file="$1" + local resource_log="$2" + local net_before_rx="$3" + local net_before_tx="$4" + local net_after_rx="$5" + local net_after_tx="$6" + local rx_delta=$((net_after_rx - net_before_rx)) + local tx_delta=$((net_after_tx - net_before_tx)) + + { + echo + echo "Resource sampling log: ${resource_log}" + summarize_resource_role "$resource_log" "server" + summarize_resource_role "$resource_log" "bench" + echo "Network bytes delta: rx_bytes=${rx_delta} tx_bytes=${tx_delta}" + } | tee -a "$log_file" +} + function start_server() { echo "Cleaning old local_data..." rm -rf local_data @@ -160,7 +280,14 @@ function run_suite() { local wait_timeout="$2" local command="$3" local log_file + local resource_log + local sampler_pid="" + local net_before_rx=0 + local net_before_tx=0 + local net_after_rx=0 + local net_after_tx=0 log_file=$(result_log_name "$suite" "$wait_timeout") + resource_log="${log_file%.log}_resources.csv" echo echo "Suite: ${suite}, poll wait timeout: ${wait_timeout}" @@ -171,15 +298,30 @@ function run_suite() { fi start_server + + if [[ "$RESOURCE_SAMPLING" == true ]]; then + echo "timestamp,role,pid,cpu_percent,rss_kb" >"$resource_log" + read -r net_before_rx net_before_tx < <(snapshot_network_bytes) + sample_process_resources "$resource_log" & + sampler_pid=$! + fi + set +e eval "$command" 2>&1 | tee "$log_file" local status=${PIPESTATUS[0]} set -e + + if [[ "$RESOURCE_SAMPLING" == true ]]; then + stop_resource_sampler "$sampler_pid" + read -r net_after_rx net_after_tx < <(snapshot_network_bytes) + append_resource_summary "$log_file" "$resource_log" "$net_before_rx" "$net_before_tx" "$net_after_rx" "$net_after_tx" + fi + stop_server echo echo "Summary from ${log_file}:" - grep -E "Results:|Total throughput|latency:" "$log_file" || true + grep -E "Results:|Total throughput|latency:|Resource summary:|Network bytes delta:" "$log_file" || true if [[ $status -ne 0 ]]; then echo "Benchmark failed with status ${status}: ${suite}, wait=${wait_timeout}" >&2 @@ -200,6 +342,13 @@ function command_for_sparse_single() { echo "${IGGY_BENCH_CMD} --message-size 1000 --messages-per-batch 1 --message-batches ${SPARSE_BATCHES} --rate-limit 10KB --poll-wait-timeout ${wait_timeout} pinned-producer-and-consumer --streams 1 --producers 1 --consumers 1 $(bench_transport) ${output_args}" } +function command_for_busy_loop_single() { + local wait_timeout="$1" + local output_args + output_args=$(common_output_args "busy_loop_single" "$wait_timeout") + echo "${IGGY_BENCH_CMD} --message-size 1000 --messages-per-batch 1 --message-batches ${BUSY_LOOP_BATCHES} --rate-limit 5KB --poll-wait-timeout ${wait_timeout} pinned-producer-and-consumer --streams 1 --producers 1 --consumers 1 $(bench_transport) ${output_args}" +} + function command_for_sparse_consumer_group() { local wait_timeout="$1" local output_args @@ -215,12 +364,16 @@ function command_for_saturated_control() { } echo "Running PollMessages wait-timeout benchmark comparison" -echo "transport=${TRANSPORT}, identifier=${IDENTIFIER}, output_dir=${OUTPUT_DIR}, quick=${QUICK}, dry_run=${DRY_RUN}" +echo "transport=${TRANSPORT}, identifier=${IDENTIFIER}, output_dir=${OUTPUT_DIR}, quick=${QUICK}, dry_run=${DRY_RUN}, resource_sampling=${RESOURCE_SAMPLING}, sample_interval=${SAMPLE_INTERVAL_SECONDS}s" for wait_timeout in "${wait_timeouts[@]}"; do run_suite "sparse_single" "$wait_timeout" "$(command_for_sparse_single "$wait_timeout")" done +for wait_timeout in "${busy_loop_wait_timeouts[@]}"; do + run_suite "busy_loop_single" "$wait_timeout" "$(command_for_busy_loop_single "$wait_timeout")" +done + for wait_timeout in "${wait_timeouts[@]}"; do run_suite "sparse_consumer_group" "$wait_timeout" "$(command_for_sparse_consumer_group "$wait_timeout")" done @@ -230,4 +383,4 @@ for wait_timeout in "0s" "100ms"; do done echo -echo "Comparison complete. Full logs and benchmark artifacts are in ${OUTPUT_DIR}." +echo "Comparison complete. Full logs, resource samples, and benchmark artifacts are in ${OUTPUT_DIR}." From 33e290419140f594f1b6ff7e588e3f0f1fff8f8e Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Sun, 5 Jul 2026 17:46:30 +0530 Subject: [PATCH 20/24] fix(bench): map websocket transport for poll comparison --- scripts/performance/run-poll-wait-timeout-comparison.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/performance/run-poll-wait-timeout-comparison.sh b/scripts/performance/run-poll-wait-timeout-comparison.sh index a322780ecc..a8baceebd9 100755 --- a/scripts/performance/run-poll-wait-timeout-comparison.sh +++ b/scripts/performance/run-poll-wait-timeout-comparison.sh @@ -140,7 +140,7 @@ function bench_transport() { echo "tcp" ;; websocket) - echo "websocket" + echo "web-socket" ;; quic) echo "quic" From 459a4b732f0a45e8a7563dfae9cc6f286139771e Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Fri, 31 Jul 2026 23:58:03 +0530 Subject: [PATCH 21/24] fix: address deferred polling review feedback --- core/bench/src/args/common.rs | 2 +- core/common/src/traits/message_client.rs | 4 +- core/integration/tests/server/cg.rs | 3 +- .../poll_messages_wait_timeout_scenario.rs | 46 ++++++- core/integration/tests/server/specific.rs | 8 +- core/server-ng/src/dispatch.rs | 18 +++ core/server-ng/src/http/wire.rs | 1 + .../messages/poll_messages_handler.rs | 46 ++++--- core/server/src/http/http_shard_wrapper.rs | 4 +- core/server/src/http/messages.rs | 2 +- core/server/src/shard/builder.rs | 1 + core/server/src/shard/mod.rs | 1 + core/server/src/shard/system/messages.rs | 19 +-- core/server/src/shard/waiters.rs | 66 +++++----- .../poll-wait-timeout-comparison.md | 119 ------------------ .../run-poll-wait-timeout-comparison.sh | 24 ++-- 16 files changed, 156 insertions(+), 208 deletions(-) delete mode 100644 scripts/performance/poll-wait-timeout-comparison.md diff --git a/core/bench/src/args/common.rs b/core/bench/src/args/common.rs index 8e77a34ce9..8047fb28f2 100644 --- a/core/bench/src/args/common.rs +++ b/core/bench/src/args/common.rs @@ -72,7 +72,7 @@ pub struct IggyBenchArgs { #[arg(long, short = 'w', default_value_t = IggyDuration::from_str(DEFAULT_WARMUP_TIME).unwrap())] pub warmup_time: IggyDuration, - /// Poll wait timeout in human readable format, e.g. "10ms", "1s" + /// Poll wait timeout, e.g. "10ms", "1s". Use "0s" to disable deferred polling. #[arg(long, default_value_t = IggyDuration::from_str(DEFAULT_POLL_WAIT_TIMEOUT).unwrap(), value_parser = IggyDuration::from_str)] pub poll_wait_timeout: IggyDuration, diff --git a/core/common/src/traits/message_client.rs b/core/common/src/traits/message_client.rs index f61a6932cb..73cd544564 100644 --- a/core/common/src/traits/message_client.rs +++ b/core/common/src/traits/message_client.rs @@ -42,7 +42,9 @@ pub trait MessageClient { ) -> Result; /// Polls messages and waits up to the timeout when no messages are immediately available. - /// A zero timeout preserves immediate polling behavior. + /// A zero timeout preserves immediate polling behavior. Non-zero timeouts require transport support. + /// When auto commit is enabled, the offset can be committed before the response is delivered. + /// Implementors should override at least one polling method with a real implementation. #[allow(clippy::too_many_arguments)] async fn poll_messages_with_timeout( &self, diff --git a/core/integration/tests/server/cg.rs b/core/integration/tests/server/cg.rs index b4138f7ae0..26aeb976af 100644 --- a/core/integration/tests/server/cg.rs +++ b/core/integration/tests/server/cg.rs @@ -15,13 +15,14 @@ // specific language governing permissions and limitations // under the License. +#[cfg(not(feature = "vsr"))] +use crate::server::scenarios::poll_messages_wait_timeout_scenario; use crate::server::scenarios::{ consumer_group_auto_commit_reconnection_scenario, consumer_group_duplicate_name_create_scenario, consumer_group_join_scenario, consumer_group_new_messages_after_restart_scenario, consumer_group_offset_cleanup_scenario, consumer_group_with_multiple_clients_polling_messages_scenario, consumer_group_with_single_client_polling_messages_scenario, - poll_messages_wait_timeout_scenario, }; use integration::iggy_harness; diff --git a/core/integration/tests/server/scenarios/poll_messages_wait_timeout_scenario.rs b/core/integration/tests/server/scenarios/poll_messages_wait_timeout_scenario.rs index 204ca78c2a..eecd3b0623 100644 --- a/core/integration/tests/server/scenarios/poll_messages_wait_timeout_scenario.rs +++ b/core/integration/tests/server/scenarios/poll_messages_wait_timeout_scenario.rs @@ -17,13 +17,17 @@ use iggy::prelude::*; use integration::harness::TestHarness; -use std::time::{Duration, Instant}; +use std::time::Duration; +#[cfg(not(feature = "vsr"))] +use std::time::Instant; +#[cfg(not(feature = "vsr"))] use tokio::time::sleep; #[cfg(not(feature = "vsr"))] use tokio::time::timeout; const STREAM_NAME: &str = "poll-wait-timeout-stream"; const TOPIC_NAME: &str = "poll-wait-timeout-topic"; +#[cfg(not(feature = "vsr"))] const CONSUMER_GROUP_NAME: &str = "poll-wait-timeout-group"; #[cfg(not(feature = "vsr"))] const PARTITION_ID: u32 = 0; @@ -116,6 +120,44 @@ pub async fn run_wake_after_append_checks(harness: &TestHarness) { .unwrap(); } +#[cfg(feature = "vsr")] +pub async fn run_server_ng_rejects_wait_timeout(harness: &TestHarness) { + let client = harness.root_client().await.expect("root client"); + client.create_stream(STREAM_NAME).await.unwrap(); + client + .create_topic( + &Identifier::named(STREAM_NAME).unwrap(), + TOPIC_NAME, + 1, + CompressionAlgorithm::default(), + None, + IggyExpiry::NeverExpire, + MaxTopicSize::ServerDefault, + ) + .await + .unwrap(); + + let result = client + .poll_messages_with_timeout( + &Identifier::named(STREAM_NAME).unwrap(), + &Identifier::named(TOPIC_NAME).unwrap(), + Some(0), + &Consumer::default(), + &PollingStrategy::offset(0), + 1, + false, + Duration::from_millis(10), + ) + .await; + + assert!(matches!(result, Err(IggyError::FeatureUnavailable))); + + client + .delete_stream(&Identifier::named(STREAM_NAME).unwrap()) + .await + .unwrap(); +} + #[cfg(not(feature = "vsr"))] async fn setup_topic(client: &IggyClient, partitions_count: u32) { client.create_stream(STREAM_NAME).await.unwrap(); @@ -308,6 +350,7 @@ async fn send_one(client: &IggyClient, id: u128, payload: &str) { .unwrap(); } +#[cfg(not(feature = "vsr"))] pub async fn run_consumer_group_checks(harness: &TestHarness) { let client = harness.root_client().await.expect("root client"); @@ -401,6 +444,7 @@ pub async fn run_consumer_group_checks(harness: &TestHarness) { .unwrap(); } +#[cfg(not(feature = "vsr"))] async fn wait_until_partition_readable(client: &IggyClient) { let deadline = Instant::now() + Duration::from_secs(5); loop { diff --git a/core/integration/tests/server/specific.rs b/core/integration/tests/server/specific.rs index d64ac4c096..86dc9d7a76 100644 --- a/core/integration/tests/server/specific.rs +++ b/core/integration/tests/server/specific.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -#[cfg(not(feature = "vsr"))] use crate::server::scenarios::poll_messages_wait_timeout_scenario; use crate::server::scenarios::{message_size_scenario, single_message_per_batch_scenario}; use crate::server::scenarios::{reconnect_after_restart_scenario, restart_offset_skip_scenario}; @@ -61,6 +60,7 @@ async fn should_handle_single_message_per_batch_with_delayed_persistence(harness single_message_per_batch_scenario::run(harness, 5).await; } +#[cfg(not(feature = "vsr"))] #[iggy_harness( test_client_transport = [Tcp, WebSocket, Quic], server(tcp.socket.override_defaults = true, tcp.socket.nodelay = true) @@ -69,6 +69,12 @@ async fn poll_messages_wait_timeout_wakes_after_append(harness: &TestHarness) { poll_messages_wait_timeout_scenario::run_wake_after_append_checks(harness).await; } +#[cfg(feature = "vsr")] +#[iggy_harness(cluster_nodes = 1, test_client_transport = [Tcp, WebSocket])] +async fn poll_messages_wait_timeout_rejected_by_vsr(harness: &TestHarness) { + poll_messages_wait_timeout_scenario::run_server_ng_rejects_wait_timeout(harness).await; +} + #[cfg(not(feature = "vsr"))] #[iggy_harness( test_client_transport = [Tcp], diff --git a/core/server-ng/src/dispatch.rs b/core/server-ng/src/dispatch.rs index b6bea8cd88..38f42a0f96 100644 --- a/core/server-ng/src/dispatch.rs +++ b/core/server-ng/src/dispatch.rs @@ -1660,6 +1660,21 @@ async fn handle_poll_messages( .await; return; }; + if wire.wait_timeout_us != 0 { + warn!( + transport_client_id, + wait_timeout_us = wire.wait_timeout_us, + "server-ng does not support deferred poll waits yet" + ); + send_non_replicated_deny( + shard, + request, + transport_client_id, + IggyError::FeatureUnavailable.as_code(), + ) + .await; + return; + } // Gate on (stream, topic) before touching the partition plane. A resolution // miss falls through to the resolve path below (empty-poll / not-found); a // denial replies status!=0 with an empty body, distinct from the empty-poll @@ -1977,6 +1992,9 @@ where MJ::Target: Journal, Header = PrepareHeader>, S: 'static, { + if wire.wait_timeout_us != 0 { + return Err(IggyError::FeatureUnavailable); + } let strategy = polling_strategy_from_wire(&wire.strategy)?; let args = PollingArgs::new(strategy, wire.count, wire.auto_commit); diff --git a/core/server-ng/src/http/wire.rs b/core/server-ng/src/http/wire.rs index e96866f80a..aa74474e07 100644 --- a/core/server-ng/src/http/wire.rs +++ b/core/server-ng/src/http/wire.rs @@ -105,6 +105,7 @@ pub(in crate::http) fn poll_wire_request( }, count: query.count, auto_commit: query.auto_commit, + wait_timeout_us: 0, }) } diff --git a/core/server/src/binary/handlers/messages/poll_messages_handler.rs b/core/server/src/binary/handlers/messages/poll_messages_handler.rs index de22934e54..925b3b420d 100644 --- a/core/server/src/binary/handlers/messages/poll_messages_handler.rs +++ b/core/server/src/binary/handlers/messages/poll_messages_handler.rs @@ -31,7 +31,9 @@ use iggy_common::{Consumer, IggyError, IggyPollMetadata, PollingStrategy}; use server_common::PooledBuffer; use server_common::sharding::IggyNamespace; use std::{rc::Rc, time::Duration}; -use tracing::{debug, trace}; +use tracing::{debug, trace, warn}; + +const MAX_POLL_WAIT_TIMEOUT: Duration = Duration::from_secs(30); pub async fn handle_poll_messages( req: PollMessagesRequest, @@ -46,7 +48,7 @@ pub async fn handle_poll_messages( let partition_id = req.partition_id; let count = req.count; let auto_commit = req.auto_commit; - let wait_timeout = Duration::from_micros(req.wait_timeout_us); + let wait_timeout = Duration::from_micros(req.wait_timeout_us).min(MAX_POLL_WAIT_TIMEOUT); debug!( "session: {session}, command: poll_messages, stream_id: {stream_id}, topic_id: {topic_id}, partition_id: {partition_id:?}" @@ -60,9 +62,9 @@ pub async fn handle_poll_messages( .poll_messages( client_id, topic, - consumer.clone(), + &consumer, partition_id, - PollingArgs::with_wait_timeout(strategy, count, auto_commit, wait_timeout), + PollingArgs::new(strategy, count, auto_commit), ) .await?; @@ -71,7 +73,16 @@ pub async fn handle_poll_messages( shard.resolve_poll_wait_namespaces(topic, &consumer, client_id, partition_id)?; let waiters = namespaces .iter() - .filter_map(|namespace| shard.register_poll_waiter(*namespace, wait_timeout)) + .filter_map(|namespace| { + let waiter = shard.register_poll_waiter(*namespace); + if waiter.is_none() { + warn!( + namespace = namespace.inner(), + "poll waiter cap reached; falling back to immediate poll" + ); + } + waiter + }) .collect::>(); if let Some((next_metadata, next_batch)) = poll_wait_namespaces( @@ -91,14 +102,15 @@ pub async fn handle_poll_messages( } if batch.is_empty() && !waiters.is_empty() { - let woke = matches!( - compio::time::timeout(wait_timeout, wait_for_any_poll_waiter(&waiters)).await, - Ok(true) - ); - drop(waiters); + let deadline = std::time::Instant::now() + wait_timeout; + while batch.is_empty() { + let remaining = deadline.saturating_duration_since(std::time::Instant::now()); + if remaining.is_zero() { + break; + } + let _ = compio::time::timeout(remaining, wait_for_any_poll_waiter(&waiters)).await; - if woke - && let Some((next_metadata, next_batch)) = poll_wait_namespaces( + if let Some((next_metadata, next_batch)) = poll_wait_namespaces( shard, client_id, topic, @@ -109,10 +121,12 @@ pub async fn handle_poll_messages( auto_commit, ) .await? - { - metadata = next_metadata; - batch = next_batch; + { + metadata = next_metadata; + batch = next_batch; + } } + drop(waiters); } } @@ -162,7 +176,7 @@ async fn poll_wait_namespaces( .poll_messages( client_id, topic, - consumer.clone(), + consumer, Some(namespace.partition_id() as u32), PollingArgs::new(*strategy, count, auto_commit), ) diff --git a/core/server/src/http/http_shard_wrapper.rs b/core/server/src/http/http_shard_wrapper.rs index 9952e4de78..e71e44b3cb 100644 --- a/core/server/src/http/http_shard_wrapper.rs +++ b/core/server/src/http/http_shard_wrapper.rs @@ -141,7 +141,7 @@ impl HttpSafeShard { user_id: u32, stream_id: Identifier, topic_id: Identifier, - consumer: Consumer, + consumer: &Consumer, maybe_partition_id: Option, args: PollingArgs, ) -> Result<(IggyPollMetadata, IggyMessagesBatchSet), IggyError> { @@ -151,7 +151,7 @@ impl HttpSafeShard { let future = SendWrapper::new(self.shard().poll_messages( client_id, topic, - consumer.clone(), + consumer, maybe_partition_id, args, )); diff --git a/core/server/src/http/messages.rs b/core/server/src/http/messages.rs index b51df4aef1..cd7cd6341d 100644 --- a/core/server/src/http/messages.rs +++ b/core/server/src/http/messages.rs @@ -70,7 +70,7 @@ async fn poll_messages( session.get_user_id(), query.0.stream_id, query.0.topic_id, - consumer, + &consumer, query.0.partition_id, PollingArgs::new(query.0.strategy, query.0.count, query.0.auto_commit), )); diff --git a/core/server/src/shard/builder.rs b/core/server/src/shard/builder.rs index f7e41c3eb3..26637514eb 100644 --- a/core/server/src/shard/builder.rs +++ b/core/server/src/shard/builder.rs @@ -183,6 +183,7 @@ impl IggyShardBuilder { poll_waiters: self .poll_waiters .unwrap_or_else(|| Arc::new(Mutex::new(PollWaiterRegistry::default()))), + poll_waiters_live: Arc::new(std::sync::atomic::AtomicUsize::new(0)), encryptor, config, _version: version, diff --git a/core/server/src/shard/mod.rs b/core/server/src/shard/mod.rs index a729f6b8bd..5d285e7707 100644 --- a/core/server/src/shard/mod.rs +++ b/core/server/src/shard/mod.rs @@ -81,6 +81,7 @@ pub struct IggyShard { pub(crate) local_partitions: RefCell, pub(crate) pending_partition_inits: RefCell>, pub(crate) poll_waiters: Arc>, + pub(crate) poll_waiters_live: Arc, pub(crate) shards_table: EternalPtr>, pub(crate) state: FileState, diff --git a/core/server/src/shard/system/messages.rs b/core/server/src/shard/system/messages.rs index 076607c98b..3fea3f8121 100644 --- a/core/server/src/shard/system/messages.rs +++ b/core/server/src/shard/system/messages.rs @@ -31,7 +31,7 @@ use iggy_common::{ }; use server_common::PooledBuffer; use server_common::sharding::IggyNamespace; -use std::{sync::atomic::Ordering, time::Duration}; +use std::sync::atomic::Ordering; use tracing::error; impl IggyShard { @@ -68,13 +68,13 @@ impl IggyShard { &self, client_id: u32, topic: ResolvedTopic, - consumer: Consumer, + consumer: &Consumer, maybe_partition_id: Option, args: PollingArgs, ) -> Result<(IggyPollMetadata, IggyMessagesBatchSet), IggyError> { let Some((consumer, partition_id)) = self.resolve_consumer_with_partition_id( topic, - &consumer, + consumer, client_id, maybe_partition_id, true, @@ -410,6 +410,8 @@ impl IggyShard { .size_of_messages_required_to_save .as_bytes_u64() as u32; + self.wake_poll_waiters(namespace); + if is_full || unsaved_messages_count_exceeded || unsaved_messages_size_exceeded { let frozen_batches = { let mut partitions = self.local_partitions.borrow_mut(); @@ -682,25 +684,14 @@ pub struct PollingArgs { pub strategy: PollingStrategy, pub count: u32, pub auto_commit: bool, - pub wait_timeout: Duration, } impl PollingArgs { pub fn new(strategy: PollingStrategy, count: u32, auto_commit: bool) -> Self { - Self::with_wait_timeout(strategy, count, auto_commit, Duration::ZERO) - } - - pub fn with_wait_timeout( - strategy: PollingStrategy, - count: u32, - auto_commit: bool, - wait_timeout: Duration, - ) -> Self { Self { strategy, count, auto_commit, - wait_timeout, } } } diff --git a/core/server/src/shard/waiters.rs b/core/server/src/shard/waiters.rs index 0af4e2b3e1..b491483f3e 100644 --- a/core/server/src/shard/waiters.rs +++ b/core/server/src/shard/waiters.rs @@ -19,7 +19,11 @@ use crate::shard::IggyShard; use ahash::AHashMap; use async_channel::{Receiver, Sender}; use server_common::sharding::IggyNamespace; -use std::{sync::Arc, sync::Mutex, time::Duration, time::Instant}; +use std::{ + sync::Arc, + sync::Mutex, + sync::atomic::{AtomicUsize, Ordering}, +}; const MAX_WAITERS_PER_NAMESPACE: usize = 1024; @@ -27,7 +31,6 @@ const MAX_WAITERS_PER_NAMESPACE: usize = 1024; struct PollWaiter { id: u64, wake_sender: Sender<()>, - deadline: Option, } #[derive(Debug, Default)] @@ -37,12 +40,7 @@ pub struct PollWaiterRegistry { } impl PollWaiterRegistry { - fn register( - &mut self, - namespace: IggyNamespace, - timeout: Duration, - ) -> Option<(u64, Receiver<()>)> { - self.prune_namespace(&namespace); + fn register(&mut self, namespace: IggyNamespace) -> Option<(u64, Receiver<()>)> { let waiters = self.waiters.entry(namespace).or_default(); if waiters.len() >= MAX_WAITERS_PER_NAMESPACE { return None; @@ -51,11 +49,7 @@ impl PollWaiterRegistry { self.next_id = self.next_id.wrapping_add(1); let id = self.next_id; let (wake_sender, wake_receiver) = async_channel::bounded(1); - waiters.push(PollWaiter { - id, - wake_sender, - deadline: Instant::now().checked_add(timeout), - }); + waiters.push(PollWaiter { id, wake_sender }); Some((id, wake_receiver)) } @@ -107,23 +101,6 @@ impl PollWaiterRegistry { self.wake_namespace(&namespace); } } - - fn prune_namespace(&mut self, namespace: &IggyNamespace) { - let now = Instant::now(); - let should_remove = if let Some(waiters) = self.waiters.get_mut(namespace) { - waiters.retain(|waiter| { - !waiter.wake_sender.is_closed() - && waiter.deadline.is_none_or(|deadline| deadline > now) - }); - waiters.is_empty() - } else { - false - }; - - if should_remove { - self.waiters.remove(namespace); - } - } } pub(crate) struct PollWaiterRegistration { @@ -131,6 +108,7 @@ pub(crate) struct PollWaiterRegistration { id: u64, receiver: Receiver<()>, registry: Arc>, + live_counter: Arc, } impl PollWaiterRegistration { @@ -143,8 +121,9 @@ impl Drop for PollWaiterRegistration { fn drop(&mut self) { self.registry .lock() - .expect("poll waiter registry poisoned") + .unwrap_or_else(|error| error.into_inner()) .remove(&self.namespace, self.id); + self.live_counter.fetch_sub(1, Ordering::Relaxed); } } @@ -152,39 +131,49 @@ impl IggyShard { pub(crate) fn register_poll_waiter( &self, namespace: IggyNamespace, - timeout: Duration, ) -> Option { let (id, receiver) = self .poll_waiters .lock() - .expect("poll waiter registry poisoned") - .register(namespace, timeout)?; + .unwrap_or_else(|error| error.into_inner()) + .register(namespace)?; + self.poll_waiters_live.fetch_add(1, Ordering::Relaxed); Some(PollWaiterRegistration { namespace, id, receiver, registry: self.poll_waiters.clone(), + live_counter: self.poll_waiters_live.clone(), }) } pub(crate) fn wake_poll_waiters(&self, namespace: &IggyNamespace) { + if self.poll_waiters_live.load(Ordering::Relaxed) == 0 { + return; + } self.poll_waiters .lock() - .expect("poll waiter registry poisoned") + .unwrap_or_else(|error| error.into_inner()) .wake_namespace(namespace); } pub(crate) fn wake_topic_poll_waiters(&self, stream_id: usize, topic_id: usize) { + if self.poll_waiters_live.load(Ordering::Relaxed) == 0 { + return; + } self.poll_waiters .lock() - .expect("poll waiter registry poisoned") + .unwrap_or_else(|error| error.into_inner()) .wake_topic(stream_id, topic_id); } pub(crate) fn wake_stream_poll_waiters(&self, stream_id: usize) { + if self.poll_waiters_live.load(Ordering::Relaxed) == 0 { + return; + } self.poll_waiters .lock() - .expect("poll waiter registry poisoned") + .unwrap_or_else(|error| error.into_inner()) .wake_stream(stream_id); } } @@ -200,7 +189,7 @@ mod tests { let (id, receiver) = registry .lock() .unwrap() - .register(namespace, Duration::from_secs(1)) + .register(namespace) .expect("waiter should register"); assert_eq!(registry.lock().unwrap().waiters[&namespace].len(), 1); @@ -210,6 +199,7 @@ mod tests { id, receiver, registry: registry.clone(), + live_counter: Arc::new(AtomicUsize::new(1)), }; drop(registration); diff --git a/scripts/performance/poll-wait-timeout-comparison.md b/scripts/performance/poll-wait-timeout-comparison.md deleted file mode 100644 index 72b7ffe0be..0000000000 --- a/scripts/performance/poll-wait-timeout-comparison.md +++ /dev/null @@ -1,119 +0,0 @@ -# Poll Wait Timeout Benchmark Comparison - -This runbook compares immediate `PollMessages` polling with deferred polling. - -## Goal - -Compare immediate polling with deferred polling on sparse, busy-loop, consumer-group, and saturated workloads. The script records `iggy-bench` throughput and latency, plus OS-level CPU/RSS samples and network byte deltas for each run. - -## What To Compare - -- Immediate polling: `--poll-wait-timeout 0s` -- Deferred polling: `--poll-wait-timeout 10ms`, `100ms`, `1s` -- Busy-loop comparison: `0s` vs `100ms` -- Transports: run `tcp` first, then repeat with `websocket` -- Sparse workloads: low producer rate, small batches -- Saturated control: large batches, no rate limit -- Resource signals: `iggy-server` CPU/RSS, `iggy-bench` CPU/RSS, network bytes before/after each run - -## Run - -Quick TCP smoke, about 2-3 minutes on a warm release build: - -```bash -scripts/performance/run-poll-wait-timeout-comparison.sh --quick --transport tcp -``` - -Full TCP run: - -```bash -scripts/performance/run-poll-wait-timeout-comparison.sh --transport tcp -``` - -Full WebSocket run: - -```bash -scripts/performance/run-poll-wait-timeout-comparison.sh --transport websocket -``` - -Use existing release binaries when you already built them elsewhere: - -```bash -scripts/performance/run-poll-wait-timeout-comparison.sh \ - --quick \ - --skip-build \ - --bench-cmd /Users/aruns/Developer/iggy/target/release/iggy-bench \ - --server-cmd /Users/aruns/Developer/iggy/target/release/iggy-server \ - --identifier aruns-mbp-pr3605 -``` - -Change OS sample interval: - -```bash -scripts/performance/run-poll-wait-timeout-comparison.sh --sample-interval 0.5 -``` - -Disable OS sampling if you only want `iggy-bench` throughput and latency: - -```bash -scripts/performance/run-poll-wait-timeout-comparison.sh --no-resource-sampling -``` - -Print commands without running: - -```bash -scripts/performance/run-poll-wait-timeout-comparison.sh --dry-run -``` - -## Workloads - -The script runs: - -- `sparse_single`: one producer and one consumer, `1` message per batch, `10KB/s` rate limit -- `busy_loop_single`: one producer and one consumer, `1` message per batch, `5KB/s` rate limit, `0s` vs `100ms` -- `sparse_consumer_group`: one producer and four consumers over four partitions, `1` message per batch, `10KB/s` rate limit -- `saturated_control`: one producer and one consumer, `1000` messages per batch, no rate limit, `0s` vs `100ms` - -## Resource Logs - -Each run writes a benchmark log and, when resource sampling is enabled, a matching resource CSV: - -```text -performance_results/poll_wait_timeout_comparison/_wait_.log -performance_results/poll_wait_timeout_comparison/_wait__resources.csv -``` - -The CSV format is: - -```text -timestamp,role,pid,cpu_percent,rss_kb -``` - -The benchmark log appends: - -- average and max CPU for `server` and `bench` -- max RSS for `server` and `bench` -- network byte delta for the run - -## Interpretation Caveats - -- Local TCP and WebSocket runs usually travel over loopback, so network bytes are useful for relative comparison, not external NIC throughput. -- OS CPU sampling is interval-based and process-level; use it to compare `0s` vs deferred polling trends, not as a microbenchmark-grade profiler. -- Busy-loop value should show up most clearly in CPU and network/request churn, especially when producer rate is sparse. -- Saturated control protects against regression when messages are already readable; deferred polling should stay close to immediate polling for throughput and p99 latency. -- Repeat quick smoke before sharing numbers; use full TCP/WebSocket runs for final PR discussion data. - -## Success Criteria - -- Deferred sparse runs should not lose messages. -- Sparse workloads should keep expected low-rate throughput while showing bounded wait-path latency. -- Busy-loop deferred runs should reduce CPU/network churn relative to `0s` immediate polling. -- Saturated control with non-zero timeout should be close to immediate polling for throughput and p99 latency. -- WebSocket comparison should behave consistently with TCP. -- Resource claims should cite the generated CSV/log data and avoid overclaiming precision. - -## Discord Update Template - -```text -I extended the #3605 benchmark runner to cover regular polling vs deferred polling in busy-loop/sparse workloads, plus OS-level CPU/RSS samples and network byte deltas per run. I’ll run TCP first, then repeat WebSocket, and share the 0s vs 100ms comparison with throughput, p99, CPU, RSS, and network deltas. -``` diff --git a/scripts/performance/run-poll-wait-timeout-comparison.sh b/scripts/performance/run-poll-wait-timeout-comparison.sh index a8baceebd9..beecffcf20 100755 --- a/scripts/performance/run-poll-wait-timeout-comparison.sh +++ b/scripts/performance/run-poll-wait-timeout-comparison.sh @@ -30,6 +30,7 @@ DRY_RUN=false QUICK=false RESOURCE_SAMPLING=true SAMPLE_INTERVAL_SECONDS=1 +RESOURCE_SAMPLER_PID="" while [[ $# -gt 0 ]]; do case "$1" in @@ -98,13 +99,19 @@ source "$(dirname "$0")/../utils.sh" source "$(dirname "$0")/utils.sh" function on_interrupt() { + stop_resource_sampler "$RESOURCE_SAMPLER_PID" on_exit_bench exit 130 } +function on_exit() { + stop_resource_sampler "$RESOURCE_SAMPLER_PID" + on_exit_bench +} + if [[ "$DRY_RUN" == false ]]; then trap on_interrupt SIGINT SIGTERM - trap on_exit_bench EXIT + trap on_exit EXIT fi if [[ "$SKIP_BUILD" == false && "$DRY_RUN" == false ]]; then @@ -207,15 +214,6 @@ function sample_process_resources() { done } -function stop_resource_sampler() { - local sampler_pid="$1" - - if [[ -n "$sampler_pid" ]]; then - kill "$sampler_pid" 2>/dev/null || true - wait "$sampler_pid" 2>/dev/null || true - fi -} - function summarize_resource_role() { local resource_log="$1" local role="$2" @@ -281,7 +279,6 @@ function run_suite() { local command="$3" local log_file local resource_log - local sampler_pid="" local net_before_rx=0 local net_before_tx=0 local net_after_rx=0 @@ -303,7 +300,7 @@ function run_suite() { echo "timestamp,role,pid,cpu_percent,rss_kb" >"$resource_log" read -r net_before_rx net_before_tx < <(snapshot_network_bytes) sample_process_resources "$resource_log" & - sampler_pid=$! + RESOURCE_SAMPLER_PID=$! fi set +e @@ -312,7 +309,8 @@ function run_suite() { set -e if [[ "$RESOURCE_SAMPLING" == true ]]; then - stop_resource_sampler "$sampler_pid" + stop_resource_sampler "$RESOURCE_SAMPLER_PID" + RESOURCE_SAMPLER_PID="" read -r net_after_rx net_after_tx < <(snapshot_network_bytes) append_resource_summary "$log_file" "$resource_log" "$net_before_rx" "$net_before_tx" "$net_after_rx" "$net_after_tx" fi From 3a21b1e41c4c2d65a3061f49ced5cc4bf5dae339 Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Sat, 1 Aug 2026 13:48:30 +0530 Subject: [PATCH 22/24] fix: update simulator poll request --- core/simulator/src/client.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/core/simulator/src/client.rs b/core/simulator/src/client.rs index f511c10006..3306cd5977 100644 --- a/core/simulator/src/client.rs +++ b/core/simulator/src/client.rs @@ -574,6 +574,7 @@ impl SimClient { strategy: WirePollingStrategy::first(), count, auto_commit: false, + wait_timeout_us: 0, } .to_bytes(); From 5e1b6a1367d6cd2e5d86883d6f591b170e1b9f1a Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Sat, 1 Aug 2026 13:58:53 +0530 Subject: [PATCH 23/24] fix: gate reconnect producer scenario for vsr --- .../server/scenarios/reconnect_after_restart_scenario.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/core/integration/tests/server/scenarios/reconnect_after_restart_scenario.rs b/core/integration/tests/server/scenarios/reconnect_after_restart_scenario.rs index bc2da26a73..2bc0e4f95e 100644 --- a/core/integration/tests/server/scenarios/reconnect_after_restart_scenario.rs +++ b/core/integration/tests/server/scenarios/reconnect_after_restart_scenario.rs @@ -20,11 +20,12 @@ use iggy::prelude::*; use iggy_common::TransportProtocol; use integration::harness::{TestBinary, TestHarness}; use std::str::FromStr; -use tokio::time::{Duration, sleep, timeout}; +use tokio::time::{Duration, sleep}; const STREAM_NAME: &str = "test-reconnect-stream"; const TOPIC_NAME: &str = "test-reconnect-topic"; +#[cfg(not(feature = "vsr"))] pub async fn run_producer(harness: &mut TestHarness) { let client = create_client(harness); Client::connect(&client).await.expect("Failed to connect"); @@ -69,7 +70,7 @@ pub async fn run_producer(harness: &mut TestHarness) { .start() .expect("Failed to start server"); - let send_result = timeout(Duration::from_secs(60), send_handle) + let send_result = tokio::time::timeout(Duration::from_secs(60), send_handle) .await .expect("Timed out waiting for send after server restart") .expect("Send task panicked"); From 8bffd34b9cb2f870959d4dea113a1cf1164436d4 Mon Sep 17 00:00:00 2001 From: Arun Singh Date: Sat, 1 Aug 2026 14:21:24 +0530 Subject: [PATCH 24/24] fix: share poll waiter live counter --- core/server/src/main.rs | 5 ++++- core/server/src/shard/builder.rs | 15 +++++++++++++-- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/core/server/src/main.rs b/core/server/src/main.rs index 2ae91c87f6..c9c89bfe69 100644 --- a/core/server/src/main.rs +++ b/core/server/src/main.rs @@ -341,6 +341,7 @@ fn main() -> Result<(), ServerError> { let client_manager: EternalPtr> = client_manager.into(); let client_manager = ClientManager::new(client_manager); let poll_waiters = Arc::new(Mutex::new(PollWaiterRegistry::default())); + let poll_waiters_live = Arc::new(std::sync::atomic::AtomicUsize::new(0)); // Populate shards_table from SharedMetadata partitions (hierarchical traversal) metadata.with_metadata(|metadata| { @@ -387,6 +388,7 @@ fn main() -> Result<(), ServerError> { ); let client_manager = client_manager.clone(); let poll_waiters = poll_waiters.clone(); + let poll_waiters_live = poll_waiters_live.clone(); let shard_metadata = metadata.clone(); // Take metadata_writer for shard 0 only @@ -435,7 +437,8 @@ fn main() -> Result<(), ServerError> { .is_follower(is_follower) .current_replica_id(replica_id) .metadata(shard_metadata) - .poll_waiters(poll_waiters); + .poll_waiters(poll_waiters) + .poll_waiters_live(poll_waiters_live); if let Some(writer) = shard_metadata_writer { builder = builder.metadata_writer(writer); diff --git a/core/server/src/shard/builder.rs b/core/server/src/shard/builder.rs index 26637514eb..6eb6bf4795 100644 --- a/core/server/src/shard/builder.rs +++ b/core/server/src/shard/builder.rs @@ -37,7 +37,10 @@ use server_common::sharding::{IggyNamespace, PartitionLocation}; use std::{ cell::{Cell, RefCell}, rc::Rc, - sync::{Arc, Mutex, atomic::AtomicBool}, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, AtomicUsize}, + }, }; #[derive(Default)] @@ -59,6 +62,7 @@ pub struct IggyShardBuilder { metadata: Option, metadata_writer: Option, poll_waiters: Option>>, + poll_waiters_live: Option>, } impl IggyShardBuilder { @@ -135,6 +139,11 @@ impl IggyShardBuilder { self } + pub fn poll_waiters_live(mut self, poll_waiters_live: Arc) -> Self { + self.poll_waiters_live = Some(poll_waiters_live); + self + } + // TODO: Too much happens in there, some of those bootstrapping logic should be moved outside. pub fn build(self) -> IggyShard { let id = self.id.unwrap(); @@ -183,7 +192,9 @@ impl IggyShardBuilder { poll_waiters: self .poll_waiters .unwrap_or_else(|| Arc::new(Mutex::new(PollWaiterRegistry::default()))), - poll_waiters_live: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + poll_waiters_live: self + .poll_waiters_live + .unwrap_or_else(|| Arc::new(AtomicUsize::new(0))), encryptor, config, _version: version,