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..8047fb28f2 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, 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, + /// 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/binary_protocol/src/requests/messages/poll_messages.rs b/core/binary_protocol/src/requests/messages/poll_messages.rs index 4c6413cfc8..82a6248f66 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,54 @@ 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; + 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], + wait_timeout_bytes.as_slice() + ); + } + #[test] fn partition_none_encodes_zero_bytes() { let req = PollMessagesRequest { @@ -185,13 +261,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 +275,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}" diff --git a/core/common/src/traits/binary_impls/messages.rs b/core/common/src/traits/binary_impls/messages.rs index 24c2a799d0..6b022d1aca 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}") @@ -181,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, @@ -189,6 +198,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 +228,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 +275,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 +315,7 @@ impl MessageClient for B { strategy, count, auto_commit, + wait_timeout_us, ) .await; } @@ -290,6 +327,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()) diff --git a/core/common/src/traits/message_client.rs b/core/common/src/traits/message_client.rs index aa332d618c..73cd544564 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,39 @@ 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. 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, + 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. diff --git a/core/integration/tests/server/cg.rs b/core/integration/tests/server/cg.rs index 697b3ddd92..26aeb976af 100644 --- a/core/integration/tests/server/cg.rs +++ b/core/integration/tests/server/cg.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::{ consumer_group_auto_commit_reconnection_scenario, consumer_group_duplicate_name_create_scenario, consumer_group_join_scenario, @@ -81,3 +83,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..eecd3b0623 --- /dev/null +++ b/core/integration/tests/server/scenarios/poll_messages_wait_timeout_scenario.rs @@ -0,0 +1,473 @@ +// 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; +#[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; + +#[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(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(); + 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(); +} + +#[cfg(not(feature = "vsr"))] +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(); +} + +#[cfg(not(feature = "vsr"))] +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/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"); diff --git a/core/integration/tests/server/specific.rs b/core/integration/tests/server/specific.rs index 0c455d045e..86dc9d7a76 100644 --- a/core/integration/tests/server/specific.rs +++ b/core/integration/tests/server/specific.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +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 +60,31 @@ 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) +)] +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], + 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/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 } 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?; diff --git a/core/server-ng/src/dispatch.rs b/core/server-ng/src/dispatch.rs index 3d0179b05e..62342cff1a 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 6007020b7b..925b3b420d 100644 --- a/core/server/src/binary/handlers/messages/poll_messages_handler.rs +++ b/core/server/src/binary/handlers/messages/poll_messages_handler.rs @@ -21,12 +21,19 @@ 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 std::rc::Rc; -use tracing::{debug, trace}; +use server_common::sharding::IggyNamespace; +use std::{rc::Rc, time::Duration}; +use tracing::{debug, trace, warn}; + +const MAX_POLL_WAIT_TIMEOUT: Duration = Duration::from_secs(30); pub async fn handle_poll_messages( req: PollMessagesRequest, @@ -41,21 +48,88 @@ 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).min(MAX_POLL_WAIT_TIMEOUT); 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 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, + partition_id, + PollingArgs::new(strategy, count, auto_commit), + ) .await?; + 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| { + 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( + 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 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 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; + } + } + drop(waiters); + } + } + let response_length = 4 + 8 + 4 + batch.size(); let response_length_bytes = response_length.to_le_bytes(); @@ -85,3 +159,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, + 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/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/main.rs b/core/server/src/main.rs index 1f0ac18d90..c9c89bfe69 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,8 @@ 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())); + let poll_waiters_live = Arc::new(std::sync::atomic::AtomicUsize::new(0)); // Populate shards_table from SharedMetadata partitions (hierarchical traversal) metadata.with_metadata(|metadata| { @@ -384,6 +387,8 @@ fn main() -> Result<(), ServerError> { state_term.clone(), ); 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 @@ -431,7 +436,9 @@ 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) + .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 817276484e..6eb6bf4795 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; @@ -37,7 +37,10 @@ use server_common::sharding::{IggyNamespace, PartitionLocation}; use std::{ cell::{Cell, RefCell}, rc::Rc, - sync::atomic::AtomicBool, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, AtomicUsize}, + }, }; #[derive(Default)] @@ -58,6 +61,8 @@ pub struct IggyShardBuilder { current_replica_id: Option, metadata: Option, metadata_writer: Option, + poll_waiters: Option>>, + poll_waiters_live: Option>, } impl IggyShardBuilder { @@ -129,6 +134,16 @@ impl IggyShardBuilder { self } + pub fn poll_waiters(mut self, poll_waiters: Arc>) -> Self { + self.poll_waiters = Some(poll_waiters); + 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(); @@ -174,6 +189,12 @@ impl IggyShardBuilder { metadata_writer: self.metadata_writer.map(RefCell::new), local_partitions, pending_partition_inits: RefCell::new(AHashSet::new()), + poll_waiters: self + .poll_waiters + .unwrap_or_else(|| Arc::new(Mutex::new(PollWaiterRegistry::default()))), + poll_waiters_live: self + .poll_waiters_live + .unwrap_or_else(|| Arc::new(AtomicUsize::new(0))), encryptor, config, _version: version, 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(()) } diff --git a/core/server/src/shard/mod.rs b/core/server/src/shard/mod.rs index 03f0bad1b1..5d285e7707 100644 --- a/core/server/src/shard/mod.rs +++ b/core/server/src/shard/mod.rs @@ -20,7 +20,9 @@ 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, @@ -41,7 +43,7 @@ use std::{ net::SocketAddr, rc::Rc, sync::{ - Arc, + Arc, Mutex, atomic::{AtomicBool, AtomicU64, Ordering}, }, time::{Duration, Instant}, @@ -56,6 +58,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 +80,8 @@ pub struct IggyShard { pub(crate) metadata_writer: Option>, 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 8770ac905c..3fea3f8121 100644 --- a/core/server/src/shard/system/messages.rs +++ b/core/server/src/shard/system/messages.rs @@ -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(); 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 new file mode 100644 index 0000000000..b491483f3e --- /dev/null +++ b/core/server/src/shard/waiters.rs @@ -0,0 +1,208 @@ +// 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::{ + sync::Arc, + sync::Mutex, + sync::atomic::{AtomicUsize, Ordering}, +}; + +const MAX_WAITERS_PER_NAMESPACE: usize = 1024; + +#[derive(Debug)] +struct PollWaiter { + id: u64, + wake_sender: Sender<()>, +} + +#[derive(Debug, Default)] +pub struct PollWaiterRegistry { + next_id: u64, + waiters: AHashMap>, +} + +impl PollWaiterRegistry { + 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; + } + + 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 }); + + Some((id, wake_receiver)) + } + + fn remove(&mut self, namespace: &IggyNamespace, id: u64) { + let should_remove = if let Some(waiters) = self.waiters.get_mut(namespace) { + waiters.retain(|waiter| waiter.id != id); + waiters.is_empty() + } else { + false + }; + + if should_remove { + 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); + } + } +} + +pub(crate) struct PollWaiterRegistration { + namespace: IggyNamespace, + id: u64, + receiver: Receiver<()>, + registry: Arc>, + live_counter: Arc, +} + +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 + .lock() + .unwrap_or_else(|error| error.into_inner()) + .remove(&self.namespace, self.id); + self.live_counter.fetch_sub(1, Ordering::Relaxed); + } +} + +impl IggyShard { + pub(crate) fn register_poll_waiter( + &self, + namespace: IggyNamespace, + ) -> Option { + let (id, receiver) = self + .poll_waiters + .lock() + .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() + .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() + .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() + .unwrap_or_else(|error| error.into_inner()) + .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) + .expect("waiter should register"); + + assert_eq!(registry.lock().unwrap().waiters[&namespace].len(), 1); + + let registration = PollWaiterRegistration { + namespace, + id, + receiver, + registry: registry.clone(), + live_counter: Arc::new(AtomicUsize::new(1)), + }; + drop(registration); + + assert!(!registry.lock().unwrap().waiters.contains_key(&namespace)); + } +} 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(); 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..beecffcf20 --- /dev/null +++ b/scripts/performance/run-poll-wait-timeout-comparison.sh @@ -0,0 +1,384 @@ +#!/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 +RESOURCE_SAMPLING=true +SAMPLE_INTERVAL_SECONDS=1 +RESOURCE_SAMPLER_PID="" + +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 + ;; + --sample-interval) + SAMPLE_INTERVAL_SECONDS="$2" + shift 2 + ;; + --no-resource-sampling) + RESOURCE_SAMPLING=false + shift + ;; + --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 + +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" + +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 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 + 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 + tcp) + echo "tcp" + ;; + websocket) + echo "web-socket" + ;; + 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 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 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 + + 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 + local resource_log + 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}" + echo "Command: ${command}" + + if [[ "$DRY_RUN" == true ]]; then + return 0 + 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" & + RESOURCE_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 "$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 + + stop_server + + echo + echo "Summary from ${log_file}:" + 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 + 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_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 + 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}, 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 + +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, resource samples, and benchmark artifacts are in ${OUTPUT_DIR}."