diff --git a/core/bench/src/actors/consumer/typed_benchmark_consumer.rs b/core/bench/src/actors/consumer/typed_benchmark_consumer.rs index ae051fceda..a7aa1fad1e 100644 --- a/core/bench/src/actors/consumer/typed_benchmark_consumer.rs +++ b/core/bench/src/actors/consumer/typed_benchmark_consumer.rs @@ -35,8 +35,8 @@ use bench_report::{ use iggy::prelude::*; pub enum TypedBenchmarkConsumer { - High(BenchmarkConsumer), - Low(BenchmarkConsumer), + High(Box>), + Low(Box>), } impl TypedBenchmarkConsumer { @@ -70,7 +70,7 @@ impl TypedBenchmarkConsumer { }; if use_high_level_api { - Self::High(BenchmarkConsumer::new( + Self::High(Box::new(BenchmarkConsumer::new( HighLevelConsumerClient::new(client_factory, config.clone()), benchmark_kind, finish_condition, @@ -78,9 +78,9 @@ impl TypedBenchmarkConsumer { moving_average_window, limit_bytes_per_second, config, - )) + ))) } else { - Self::Low(BenchmarkConsumer::new( + Self::Low(Box::new(BenchmarkConsumer::new( LowLevelConsumerClient::new(client_factory, config.clone()), benchmark_kind, finish_condition, @@ -88,7 +88,7 @@ impl TypedBenchmarkConsumer { moving_average_window, limit_bytes_per_second, config, - )) + ))) } } diff --git a/core/sdk/src/clients/consumer.rs b/core/sdk/src/clients/consumer.rs index ea5e40440e..b9d23a684f 100644 --- a/core/sdk/src/clients/consumer.rs +++ b/core/sdk/src/clients/consumer.rs @@ -30,6 +30,7 @@ use iggy_common::{ PollingStrategy, }; use std::collections::VecDeque; +use std::fmt::{self, Debug, Formatter}; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -92,6 +93,174 @@ pub enum AutoCommitAfter { ConsumingEveryNthMessage(u32), } +/// A cheap, cloneable view of the state shared with an [`IggyConsumer`]. +/// +/// Consuming borrows the consumer as `&mut` for the whole run, so reading its getters or +/// committing an offset concurrently means sharing it behind a lock and then waiting on +/// that lock. This view carries the same shared state and needs neither. +/// +/// Every getter is an independent load rather than part of one snapshot, so the partition +/// ID can already have moved on by the time an offset is read for it. +#[derive(Clone)] +pub struct IggyConsumerState { + client: IggyRwLock, + consumer: Arc, + stream_id: Arc, + topic_id: Arc, + is_consumer_group: bool, + allow_replay: bool, + current_partition_id: Arc, + last_consumed_offsets: Arc>, + last_stored_offsets: Arc>, +} + +impl Debug for IggyConsumerState { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("IggyConsumerState") + .field("consumer", &self.consumer) + .field("stream_id", &self.stream_id) + .field("topic_id", &self.topic_id) + .field("is_consumer_group", &self.is_consumer_group) + .field("allow_replay", &self.allow_replay) + .field("current_partition_id", &self.partition_id()) + .finish_non_exhaustive() + } +} + +impl IggyConsumerState { + fn new( + client: IggyRwLock, + consumer: Arc, + stream_id: Arc, + topic_id: Arc, + is_consumer_group: bool, + allow_replay: bool, + ) -> Self { + Self { + client, + consumer, + stream_id, + topic_id, + is_consumer_group, + allow_replay, + current_partition_id: Arc::new(AtomicU32::new(0)), + last_consumed_offsets: Arc::new(DashMap::new()), + last_stored_offsets: Arc::new(DashMap::new()), + } + } + + /// Returns the current partition ID of the consumer. + pub fn partition_id(&self) -> u32 { + self.current_partition_id.load(ORDERING) + } + + /// Retrieves the last consumed offset for the specified partition ID, or `None` while + /// the partition is still untracked. Polling seeds an entry the first time it sees a + /// partition, so `Some(0)` also covers "seen, nothing consumed yet". + /// To get the current partition ID use `partition_id()` + pub fn get_last_consumed_offset(&self, partition_id: u32) -> Option { + let offset = self.last_consumed_offsets.get(&partition_id)?; + Some(offset.load(ORDERING)) + } + + /// Retrieves the last stored offset (on the server) for the specified partition ID, or + /// `None` while the partition is still untracked. Storing seeds an entry the first time + /// it sees a partition, so `Some(0)` also covers "seen, nothing stored yet". + /// To get the current partition ID use `partition_id()` + pub fn get_last_stored_offset(&self, partition_id: u32) -> Option { + let offset = self.last_stored_offsets.get(&partition_id)?; + Some(offset.load(ORDERING)) + } + + /// Stores the consumer offset on the server either for the current partition or the provided partition ID. + pub async fn store_offset( + &self, + offset: u64, + partition_id: Option, + ) -> Result<(), IggyError> { + let partition_id = partition_id.unwrap_or_else(|| self.partition_id()); + self.store_consumer_offset(partition_id, offset, self.allow_replay) + .await + } + + /// Deletes the consumer offset on the server either for the current partition or the provided partition ID. + pub async fn delete_offset(&self, mut partition_id: Option) -> Result<(), IggyError> { + // `None` is only resolved server-side for consumer groups. For a standalone consumer + // explicitly assign the current partition_id. + if partition_id.is_none() && !self.is_consumer_group { + partition_id = Some(self.partition_id()); + } + let client = self.client.read().await; + client + .delete_consumer_offset( + &self.consumer, + &self.stream_id, + &self.topic_id, + partition_id, + ) + .await + } + + async fn store_consumer_offset( + &self, + partition_id: u32, + offset: u64, + allow_replay: bool, + ) -> Result<(), IggyError> { + let consumer = &self.consumer; + let stream_id = &self.stream_id; + let topic_id = &self.topic_id; + trace!( + "Storing offset: {offset} for consumer: {consumer}, partition ID: {partition_id}, topic: {topic_id}, stream: {stream_id}..." + ); + let stored_offset; + if let Some(offset_entry) = self.last_stored_offsets.get(&partition_id) { + stored_offset = offset_entry.load(ORDERING); + } else { + stored_offset = 0; + self.last_stored_offsets + .insert(partition_id, AtomicU64::new(0)); + } + + if !allow_replay && (offset <= stored_offset && offset >= 1) { + trace!( + "Offset: {offset} is less than or equal to the last stored offset: {stored_offset} for consumer: {consumer}, partition ID: {partition_id}, topic: {topic_id}, stream: {stream_id}. Skipping storing the offset." + ); + return Ok(()); + } + + let client = self.client.read().await; + if let Err(error) = client + .store_consumer_offset(consumer, stream_id, topic_id, Some(partition_id), offset) + .await + { + error!( + "Failed to store offset: {offset} for consumer: {consumer}, partition ID: {partition_id}, topic: {topic_id}, stream: {stream_id}. {error}" + ); + return Err(error); + } + trace!( + "Stored offset: {offset} for consumer: {consumer}, partition ID: {partition_id}, topic: {topic_id}, stream: {stream_id}." + ); + if let Some(last_offset_entry) = self.last_stored_offsets.get(&partition_id) { + last_offset_entry.store(offset, ORDERING); + } else { + self.last_stored_offsets + .insert(partition_id, AtomicU64::new(offset)); + } + Ok(()) + } + + /// Snapshots the last consumed offset of every tracked partition. Collecting up front + /// releases the map guards, which must not be held across the store round trip. + fn last_consumed_offsets(&self) -> Vec<(u32, u64)> { + self.last_consumed_offsets + .iter() + .map(|entry| (*entry.key(), entry.load(ORDERING))) + .collect() + } +} + // SAFETY: IggyConsumer is Sync because: // 1. The only non-Sync field is `poll_future: Option` // 2. `poll_future` is only accessed through `poll_next()` which requires `Pin<&mut Self>` @@ -120,8 +289,7 @@ pub struct IggyConsumer { auto_commit_after_polling: bool, auto_join_consumer_group: bool, create_consumer_group_if_not_exists: bool, - last_stored_offsets: Arc>, - last_consumed_offsets: Arc>, + state: IggyConsumerState, current_offsets: Arc>, poll_future: Option, buffered_messages: VecDeque, @@ -134,7 +302,6 @@ pub struct IggyConsumer { store_offset_after_all_messages: bool, store_after_every_nth_message: u64, last_polled_at: Arc, - current_partition_id: Arc, reconnection_retry_interval: NonZeroIggyDuration, init_retries: Option, init_retry_interval: NonZeroIggyDuration, @@ -165,22 +332,33 @@ impl IggyConsumer { offset_drain_timeout: IggyDuration, ) -> Self { let (store_offset_sender, _) = flume::unbounded(); + let is_consumer_group = consumer.kind == ConsumerKind::ConsumerGroup; + let consumer = Arc::new(consumer); + let stream_id = Arc::new(stream_id); + let topic_id = Arc::new(topic_id); + let state = IggyConsumerState::new( + client.clone(), + consumer.clone(), + stream_id.clone(), + topic_id.clone(), + is_consumer_group, + allow_replay, + ); Self { initialized: false, shutdown: Arc::new(AtomicBool::new(false)), - is_consumer_group: consumer.kind == ConsumerKind::ConsumerGroup, + is_consumer_group, joined_consumer_group: Arc::new(AtomicBool::new(false)), can_poll: Arc::new(AtomicBool::new(true)), client, consumer_name, - consumer: Arc::new(consumer), - stream_id: Arc::new(stream_id), - topic_id: Arc::new(topic_id), + consumer, + stream_id, + topic_id, partition_id, polling_strategy, poll_interval_micros: polling_interval.map_or(0, |interval| interval.as_micros()), - last_stored_offsets: Arc::new(DashMap::new()), - last_consumed_offsets: Arc::new(DashMap::new()), + state, current_offsets: Arc::new(DashMap::new()), poll_future: None, batch_length, @@ -216,7 +394,6 @@ impl IggyConsumer { _ => 0, }, last_polled_at: Arc::new(AtomicU64::new(0)), - current_partition_id: Arc::new(AtomicU32::new(0)), reconnection_retry_interval, init_retries, init_retry_interval, @@ -246,7 +423,12 @@ impl IggyConsumer { /// Returns the current partition ID of the consumer. pub fn partition_id(&self) -> u32 { - self.current_partition_id.load(ORDERING) + self.state.partition_id() + } + + /// Returns a view of the consumer state that can be read without exclusive access. + pub fn state(&self) -> IggyConsumerState { + self.state.clone() } /// Stores the consumer offset on the server either for the current partition or the provided partition ID. @@ -255,54 +437,24 @@ impl IggyConsumer { offset: u64, partition_id: Option, ) -> Result<(), IggyError> { - let partition_id = if let Some(partition_id) = partition_id { - partition_id - } else { - self.current_partition_id.load(ORDERING) - }; - Self::store_consumer_offset( - &self.client, - &self.consumer, - &self.stream_id, - &self.topic_id, - partition_id, - offset, - &self.last_stored_offsets, - self.allow_replay, - ) - .await + self.state.store_offset(offset, partition_id).await } /// Retrieves the last consumed offset for the specified partition ID. /// To get the current partition ID use `partition_id()` pub fn get_last_consumed_offset(&self, partition_id: u32) -> Option { - let offset = self.last_consumed_offsets.get(&partition_id)?; - Some(offset.load(ORDERING)) + self.state.get_last_consumed_offset(partition_id) } /// Deletes the consumer offset on the server either for the current partition or the provided partition ID. - pub async fn delete_offset(&self, mut partition_id: Option) -> Result<(), IggyError> { - // `None` is only resolved server-side for consumer groups. For a standalone consumer - // explicitly assign the current partition_id. - if partition_id.is_none() && !self.is_consumer_group { - partition_id = Some(self.current_partition_id.load(ORDERING)); - } - let client = self.client.read().await; - client - .delete_consumer_offset( - &self.consumer, - &self.stream_id, - &self.topic_id, - partition_id, - ) - .await + pub async fn delete_offset(&self, partition_id: Option) -> Result<(), IggyError> { + self.state.delete_offset(partition_id).await } /// Retrieves the last stored offset (on the server) for the specified partition ID. /// To get the current partition ID use `partition_id()` pub fn get_last_stored_offset(&self, partition_id: u32) -> Option { - let offset = self.last_stored_offsets.get(&partition_id)?; - Some(offset.load(ORDERING)) + self.state.get_last_stored_offset(partition_id) } /// Initializes the consumer by subscribing to diagnostic events, initializing the consumer group if needed, storing the offsets in the background etc. @@ -397,30 +549,19 @@ impl IggyConsumer { _ => {} } - let client = self.client.clone(); - let consumer = self.consumer.clone(); - let stream_id = self.stream_id.clone(); - let topic_id = self.topic_id.clone(); - let last_stored_offsets = self.last_stored_offsets.clone(); + let state = self.state.clone(); let (store_offset_sender, store_offset_receiver) = flume::unbounded(); self.store_offset_sender = store_offset_sender; self.store_offset_task = Some(tokio::spawn(async move { while let Ok((partition_id, offset)) = store_offset_receiver.recv_async().await { trace!( - "Received offset to store: {offset}, partition ID: {partition_id}, stream: {stream_id}, topic: {topic_id}" + "Received offset to store: {offset}, partition ID: {partition_id}, stream: {}, topic: {}", + state.stream_id, state.topic_id ); - _ = Self::store_consumer_offset( - &client, - &consumer, - &stream_id, - &topic_id, - partition_id, - offset, - &last_stored_offsets, - false, - ) - .await + _ = state + .store_consumer_offset(partition_id, offset, false) + .await } })); @@ -432,63 +573,8 @@ impl IggyConsumer { Ok(()) } - #[allow(clippy::too_many_arguments)] - async fn store_consumer_offset( - client: &IggyRwLock, - consumer: &Consumer, - stream_id: &Identifier, - topic_id: &Identifier, - partition_id: u32, - offset: u64, - last_stored_offsets: &DashMap, - allow_replay: bool, - ) -> Result<(), IggyError> { - trace!( - "Storing offset: {offset} for consumer: {consumer}, partition ID: {partition_id}, topic: {topic_id}, stream: {stream_id}..." - ); - let stored_offset; - if let Some(offset_entry) = last_stored_offsets.get(&partition_id) { - stored_offset = offset_entry.load(ORDERING); - } else { - stored_offset = 0; - last_stored_offsets.insert(partition_id, AtomicU64::new(0)); - } - - if !allow_replay && (offset <= stored_offset && offset >= 1) { - trace!( - "Offset: {offset} is less than or equal to the last stored offset: {stored_offset} for consumer: {consumer}, partition ID: {partition_id}, topic: {topic_id}, stream: {stream_id}. Skipping storing the offset." - ); - return Ok(()); - } - - let client = client.read().await; - if let Err(error) = client - .store_consumer_offset(consumer, stream_id, topic_id, Some(partition_id), offset) - .await - { - error!( - "Failed to store offset: {offset} for consumer: {consumer}, partition ID: {partition_id}, topic: {topic_id}, stream: {stream_id}. {error}" - ); - return Err(error); - } - trace!( - "Stored offset: {offset} for consumer: {consumer}, partition ID: {partition_id}, topic: {topic_id}, stream: {stream_id}." - ); - if let Some(last_offset_entry) = last_stored_offsets.get(&partition_id) { - last_offset_entry.store(offset, ORDERING); - } else { - last_stored_offsets.insert(partition_id, AtomicU64::new(offset)); - } - Ok(()) - } - fn store_offsets_in_background(&self, interval: NonZeroIggyDuration) -> JoinHandle<()> { - let client = self.client.clone(); - let consumer = self.consumer.clone(); - let stream_id = self.stream_id.clone(); - let topic_id = self.topic_id.clone(); - let last_consumed_offsets = self.last_consumed_offsets.clone(); - let last_stored_offsets = self.last_stored_offsets.clone(); + let state = self.state.clone(); let shutdown = self.shutdown.clone(); let notify = self.background_commit_notify.clone(); tokio::spawn(async move { @@ -504,20 +590,10 @@ impl IggyConsumer { trace!("Shutdown signal received, stopping background offset storage"); break; } - for entry in last_consumed_offsets.iter() { - let partition_id = *entry.key(); - let consumed_offset = entry.load(ORDERING); - _ = Self::store_consumer_offset( - &client, - &consumer, - &stream_id, - &topic_id, - partition_id, - consumed_offset, - &last_stored_offsets, - false, - ) - .await; + for (partition_id, consumed_offset) in state.last_consumed_offsets() { + _ = state + .store_consumer_offset(partition_id, consumed_offset, false) + .await; } } }) @@ -680,8 +756,8 @@ impl IggyConsumer { let last_polled_at = self.last_polled_at.clone(); let can_poll = self.can_poll.clone(); let retry_interval = self.reconnection_retry_interval; - let last_stored_offset = self.last_stored_offsets.clone(); - let last_consumed_offset = self.last_consumed_offsets.clone(); + let last_stored_offset = self.state.last_stored_offsets.clone(); + let last_consumed_offset = self.state.last_consumed_offsets.clone(); let allow_replay = self.allow_replay; let is_consumer_group = self.is_consumer_group; let auto_join_consumer_group = self.auto_join_consumer_group; @@ -979,15 +1055,16 @@ impl Stream for IggyConsumer { type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let partition_id = self.current_partition_id.load(ORDERING); + let partition_id = self.state.partition_id(); if let Some(message) = self.buffered_messages.pop_front() { { if let Some(last_consumed_offset_entry) = - self.last_consumed_offsets.get(&partition_id) + self.state.last_consumed_offsets.get(&partition_id) { last_consumed_offset_entry.store(message.header.offset, ORDERING); } else { - self.last_consumed_offsets + self.state + .last_consumed_offsets .insert(partition_id, AtomicU64::new(message.header.offset)); } @@ -1032,7 +1109,9 @@ impl Stream for IggyConsumer { match future.poll_unpin(cx) { Poll::Ready(Ok(mut polled_messages)) => { let partition_id = polled_messages.partition_id; - self.current_partition_id.store(partition_id, ORDERING); + self.state + .current_partition_id + .store(partition_id, ORDERING); if polled_messages.messages.is_empty() { self.poll_future = Some(Box::pin(self.create_poll_messages_future())); } else { @@ -1088,11 +1167,12 @@ impl Stream for IggyConsumer { } if let Some(last_consumed_offset_entry) = - self.last_consumed_offsets.get(&partition_id) + self.state.last_consumed_offsets.get(&partition_id) { last_consumed_offset_entry.store(message.header.offset, ORDERING); } else { - self.last_consumed_offsets + self.state + .last_consumed_offsets .insert(partition_id, AtomicU64::new(message.header.offset)); } @@ -1171,32 +1251,18 @@ impl IggyConsumer { ); } - for entry in self.last_consumed_offsets.iter() { - let partition_id = *entry.key(); - let consumed_offset = entry.load(ORDERING); - - let stored_offset = self - .last_stored_offsets - .get(&partition_id) - .map(|e| e.load(ORDERING)) - .unwrap_or(0); + for (partition_id, consumed_offset) in self.state.last_consumed_offsets() { + let stored_offset = self.state.get_last_stored_offset(partition_id).unwrap_or(0); if consumed_offset > stored_offset { trace!( "Flushing final offset: {consumed_offset} for partition: {partition_id}, stream: {}, topic: {}", self.stream_id, self.topic_id ); - let _ = Self::store_consumer_offset( - &self.client, - &self.consumer, - &self.stream_id, - &self.topic_id, - partition_id, - consumed_offset, - &self.last_stored_offsets, - self.allow_replay, - ) - .await; + let _ = self + .state + .store_consumer_offset(partition_id, consumed_offset, self.allow_replay) + .await; } } diff --git a/core/sdk/src/prelude.rs b/core/sdk/src/prelude.rs index 326432ce3f..81f7d8c16d 100644 --- a/core/sdk/src/prelude.rs +++ b/core/sdk/src/prelude.rs @@ -33,7 +33,7 @@ pub use crate::client_wrappers::connection_info::ConnectionInfo; pub use crate::clients::client::IggyClient; pub use crate::clients::client_builder::IggyClientBuilder; pub use crate::clients::consumer::{ - AutoCommit, AutoCommitAfter, AutoCommitWhen, IggyConsumer, ReceivedMessage, + AutoCommit, AutoCommitAfter, AutoCommitWhen, IggyConsumer, IggyConsumerState, ReceivedMessage, }; pub use crate::clients::consumer_builder::IggyConsumerBuilder; pub use crate::clients::producer::IggyProducer; diff --git a/foreign/python/apache_iggy.pyi b/foreign/python/apache_iggy.pyi index e93a8ebfab..80137b39fa 100644 --- a/foreign/python/apache_iggy.pyi +++ b/foreign/python/apache_iggy.pyi @@ -1412,11 +1412,15 @@ class IggyConsumer: self, partition_id: builtins.int ) -> builtins.int | None: r""" - Get the last consumed offset or `None` if no offset has been consumed yet. + Get the last consumed offset for the given partition, or `None` while that partition + is untracked. Polling starts tracking a partition at `0`, so `0` also means + "seen, nothing consumed yet". """ def get_last_stored_offset(self, partition_id: builtins.int) -> builtins.int | None: r""" - Get the last stored offset or `None` if no offset has been stored yet. + Get the last stored offset for the given partition, or `None` while that partition is + untracked. Polling starts tracking a partition at `0`, so `0` also means + "seen, nothing stored yet", including under `AutoCommit.Disabled()`. """ def name(self) -> builtins.str: r""" @@ -1428,11 +1432,11 @@ class IggyConsumer: """ def stream(self) -> builtins.str | builtins.int: r""" - Gets the name of the stream this consumer group is configured for. + Gets the identifier of the stream this consumer group is configured for. """ def topic(self) -> builtins.str | builtins.int: r""" - Gets the name of the topic this consumer group is configured for. + Gets the identifier of the topic this consumer group is configured for. """ def store_offset( self, offset: builtins.int, partition_id: builtins.int | None diff --git a/foreign/python/pyproject.toml b/foreign/python/pyproject.toml index 00903f9267..b7c0d2f9f0 100644 --- a/foreign/python/pyproject.toml +++ b/foreign/python/pyproject.toml @@ -133,6 +133,10 @@ all = [ asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "session" timeout = 30 +# A test that pins the GIL leaves neither pytest-timeout nor asyncio able to fire. +# The faulthandler watchdog needs no GIL, so it dumps every thread and aborts the run. +faulthandler_timeout = 60 +faulthandler_exit_on_timeout = true testpaths = ["tests"] python_files = ["test_*.py"] python_classes = ["Test*"] diff --git a/foreign/python/src/client.rs b/foreign/python/src/client.rs index 00f4122f3b..f669a87484 100644 --- a/foreign/python/src/client.rs +++ b/foreign/python/src/client.rs @@ -1201,8 +1201,16 @@ impl IggyClient { .init() .await .map_err(|e| PyErr::new::(e.to_string()))?; + let state = consumer.state(); + let name = consumer.name().to_string(); + let stream = PyIdentifier::try_from(consumer.stream())?; + let topic = PyIdentifier::try_from(consumer.topic())?; Ok(IggyConsumer { inner: Arc::new(Mutex::new(consumer)), + state, + name, + stream, + topic, }) }) } diff --git a/foreign/python/src/consumer.rs b/foreign/python/src/consumer.rs index 74f521ad0a..e313a3be9c 100644 --- a/foreign/python/src/consumer.rs +++ b/foreign/python/src/consumer.rs @@ -24,7 +24,7 @@ use iggy::prelude::{ AutoCommitWhen as RustAutoCommitWhen, Consumer as RustConsumer, ConsumerGroup as RustConsumerGroup, ConsumerGroupDetails as RustConsumerGroupDetails, ConsumerGroupMember as RustConsumerGroupMember, Identifier, IggyConsumer as RustIggyConsumer, - IggyError, NonZeroIggyDuration, ReceivedMessage, + IggyConsumerState as RustIggyConsumerState, IggyError, NonZeroIggyDuration, ReceivedMessage, }; use pyo3::exceptions::PyStopAsyncIteration; use pyo3::types::PyDelta; @@ -44,51 +44,55 @@ use crate::receive_message::ReceiveMessage; /// A Python class representing the Iggy consumer. /// It provides asynchronous functionality through the contained runtime. +// `inner` stays locked for the whole duration of a consumption run, so everything that can +// be served from `state` or from a snapshot must not touch it. #[gen_stub_pyclass] #[pyclass] pub struct IggyConsumer { pub(crate) inner: Arc>, + pub(crate) state: RustIggyConsumerState, + pub(crate) name: String, + pub(crate) stream: PyIdentifier, + pub(crate) topic: PyIdentifier, } #[gen_stub_pymethods] #[pymethods] impl IggyConsumer { - /// Get the last consumed offset or `None` if no offset has been consumed yet. + /// Get the last consumed offset for the given partition, or `None` while that partition + /// is untracked. Polling starts tracking a partition at `0`, so `0` also means + /// "seen, nothing consumed yet". #[gen_stub(override_return_type(type_repr = "builtins.int | None"))] fn get_last_consumed_offset(&self, partition_id: u32) -> Option { - self.inner - .blocking_lock() - .get_last_consumed_offset(partition_id) + self.state.get_last_consumed_offset(partition_id) } - /// Get the last stored offset or `None` if no offset has been stored yet. + /// Get the last stored offset for the given partition, or `None` while that partition is + /// untracked. Polling starts tracking a partition at `0`, so `0` also means + /// "seen, nothing stored yet", including under `AutoCommit.Disabled()`. #[gen_stub(override_return_type(type_repr = "builtins.int | None"))] fn get_last_stored_offset(&self, partition_id: u32) -> Option { - self.inner - .blocking_lock() - .get_last_stored_offset(partition_id) + self.state.get_last_stored_offset(partition_id) } /// Gets the name of the consumer group. - fn name(&self) -> String { - self.inner.blocking_lock().name().to_string() + fn name(&self) -> &str { + &self.name } /// Gets the current partition id or `0` if no messages have been polled yet. fn partition_id(&self) -> u32 { - self.inner.blocking_lock().partition_id() + self.state.partition_id() } - /// Gets the name of the stream this consumer group is configured for. - fn stream(&self) -> PyResult { - let guard = self.inner.blocking_lock(); - PyIdentifier::try_from(guard.stream()) + /// Gets the identifier of the stream this consumer group is configured for. + fn stream(&self) -> PyIdentifier { + self.stream.clone() } - /// Gets the name of the topic this consumer group is configured for. - fn topic(&self) -> PyResult { - let guard = self.inner.blocking_lock(); - PyIdentifier::try_from(guard.topic()) + /// Gets the identifier of the topic this consumer group is configured for. + fn topic(&self) -> PyIdentifier { + self.topic.clone() } /// Stores the provided offset for the provided partition id or if none is specified @@ -101,11 +105,9 @@ impl IggyConsumer { offset: u64, #[gen_stub(override_type(type_repr = "builtins.int | None"))] partition_id: Option, ) -> PyResult> { - let inner = self.inner.clone(); + let state = self.state.clone(); future_into_py(py, async move { - inner - .lock() - .await + state .store_offset(offset, partition_id) .await .map_err(|e| PyErr::new::(e.to_string())) @@ -121,11 +123,9 @@ impl IggyConsumer { py: Python<'a>, #[gen_stub(override_type(type_repr = "builtins.int | None"))] partition_id: Option, ) -> PyResult> { - let inner = self.inner.clone(); + let state = self.state.clone(); future_into_py(py, async move { - inner - .lock() - .await + state .delete_offset(partition_id) .await .map_err(|e| PyErr::new::(e.to_string())) @@ -175,36 +175,32 @@ impl IggyConsumer { let mut inner = inner.lock().await; Ok(inner.consume_messages(&consumer, shutdown_rx).await) })); - let consume_result; - - if let Some(shutdown_event) = shutdown_event { - let task_locals = Python::attach(pyo3_async_runtimes::tokio::get_current_locals)?; - async fn shutdown_impl( - shutdown_event: Py, - shutdown_tx: Sender<()>, - ) -> PyResult<()> { - Python::attach(|py| { - into_future(shutdown_event.bind(py).as_any().call_method0("wait")?) - })? - .await?; - shutdown_tx.send(()).map_err(|_| { - PyErr::new::( - "Failed to signal shutdown", - ) - })?; - Ok(()) + let handle_shutdown = shutdown_event + .map(|shutdown_event| -> PyResult>> { + let task_locals = + Python::attach(pyo3_async_runtimes::tokio::get_current_locals)?; + Ok(get_runtime().spawn(scope( + task_locals, + wait_for_shutdown(shutdown_event, shutdown_tx), + ))) + }) + .transpose()?; + + let consume_result = handle_consume.await; + + if let Some(handle_shutdown) = handle_shutdown { + // Consuming can also end on its own, and the shutdown task would then park on + // `Event.wait()` forever. + handle_shutdown.abort(); + match handle_shutdown.await { + Ok(shutdown_result) => shutdown_result?, + Err(error) if error.is_cancelled() => {} + Err(error) => { + return Err(PyErr::new::( + error.to_string(), + )); + } } - let handle_shutdown: JoinHandle> = get_runtime().spawn(scope( - task_locals, - shutdown_impl(shutdown_event, shutdown_tx), - )); - let shutdown_result; - (consume_result, shutdown_result) = tokio::join!(handle_consume, handle_shutdown); - shutdown_result.map_err(|e| { - PyErr::new::(e.to_string()) - })??; - } else { - consume_result = handle_consume.await; } consume_result @@ -215,6 +211,15 @@ impl IggyConsumer { } } +async fn wait_for_shutdown(shutdown_event: Py, shutdown_tx: Sender<()>) -> PyResult<()> { + Python::attach(|py| into_future(shutdown_event.bind(py).as_any().call_method0("wait")?))? + .await?; + // A closed receiver only means consuming has already stopped, so there is nothing left + // to signal and the result of the run is the one worth reporting. + let _ = shutdown_tx.send(()); + Ok(()) +} + /// The consumer polling the messages. It selects both the consumer kind and the /// identifier the server keys the stored offset on. #[derive(Clone)] @@ -407,8 +412,7 @@ struct PyCallbackConsumer { impl MessageConsumer for PyCallbackConsumer { async fn consume(&self, received: ReceivedMessage) -> Result<(), IggyError> { let callback = self.callback.clone(); - let task_locals = self.task_locals.clone().lock_owned().await; - let task_locals = task_locals.clone(); + let task_locals = self.task_locals.lock().await.clone(); let message = ReceiveMessage { inner: received.message, partition_id: received.partition_id, diff --git a/foreign/python/tests/test_consumer_group.py b/foreign/python/tests/test_consumer_group.py index f8adf82621..31f6f06237 100644 --- a/foreign/python/tests/test_consumer_group.py +++ b/foreign/python/tests/test_consumer_group.py @@ -856,6 +856,124 @@ async def test_consumer_group_metadata(self, iggy_client: IggyClient, unique_nam assert consumer.get_last_consumed_offset(partition_id) is None assert consumer.get_last_stored_offset(partition_id) is None + @pytest.mark.asyncio + async def test_consumer_group_metadata_while_consuming( + self, iggy_client: IggyClient, unique_name + ): + """Test that metadata can be read while a consumption run is in progress.""" + consumer_name = unique_name() + stream_name = unique_name() + topic_name = unique_name() + partition_id = 0 + message = f"Metadata test - {unique_name()}" + received_messages = [] + consuming = asyncio.Event() + shutdown_event = asyncio.Event() + + await iggy_client.create_stream(stream_name) + await iggy_client.create_topic( + stream=stream_name, + name=topic_name, + partitions_count=1, + ) + + consumer = await iggy_client.consumer_group( + consumer_name, + stream_name, + topic_name, + partition_id, + PollingStrategy.First(), + 10, + auto_commit=AutoCommit.Disabled(), + poll_interval=timedelta(milliseconds=25), + ) + + async def take(received: ReceiveMessage) -> None: + received_messages.append(received) + consuming.set() + + await iggy_client.send_messages( + stream_name, + topic_name, + partition_id, + [Message(message)], + ) + + consume = consumer.consume_messages(take, shutdown_event) + try: + await asyncio.wait_for(consuming.wait(), timeout=10) + + current_partition_id = consumer.partition_id() + assert consumer.name() == consumer_name + assert consumer.stream() == stream_name + assert consumer.topic() == topic_name + assert ( + consumer.get_last_consumed_offset(current_partition_id) + == received_messages[-1].offset() + ) + assert consumer.get_last_stored_offset(current_partition_id) == 0 + finally: + shutdown_event.set() + await consume + + @pytest.mark.asyncio + async def test_consumer_group_manual_commit_while_consuming( + self, iggy_client: IggyClient, unique_name + ): + """Test that offsets can be committed from inside a consumption callback.""" + consumer_name = unique_name() + stream_name = unique_name() + topic_name = unique_name() + partition_id = 0 + messages = [f"Manual commit {i} - {unique_name()}" for i in range(2)] + received_messages = [] + stored_offsets = [] + committed = asyncio.Event() + shutdown_event = asyncio.Event() + + await iggy_client.create_stream(stream_name) + await iggy_client.create_topic( + stream=stream_name, + name=topic_name, + partitions_count=1, + ) + + consumer = await iggy_client.consumer_group( + consumer_name, + stream_name, + topic_name, + partition_id, + PollingStrategy.First(), + 10, + auto_commit=AutoCommit.Disabled(), + poll_interval=timedelta(milliseconds=25), + ) + + async def take(received: ReceiveMessage) -> None: + received_messages.append(received) + await consumer.store_offset(received.offset(), None) + stored_offsets.append(consumer.get_last_stored_offset(partition_id)) + if len(received_messages) == len(messages): + await consumer.delete_offset(partition_id) + committed.set() + + await iggy_client.send_messages( + stream_name, + topic_name, + partition_id, + [Message(message) for message in messages], + ) + + consume = consumer.consume_messages(take, shutdown_event) + try: + await asyncio.wait_for(committed.wait(), timeout=10) + assert stored_offsets == [ + received.offset() for received in received_messages + ] + finally: + shutdown_event.set() + await consume + @pytest.mark.asyncio async def test_get_last_consumed_offset_updates_as_messages_are_consumed( self, iggy_client: IggyClient, unique_name