feat(sdk): add Consumer Type extension (#1516)

Author: marvin-hansen

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "iggy"
-version = "0.6.102"
+version = "0.6.110"
 description = "Iggy is the persistent message streaming platform written in Rust, supporting QUIC, TCP and HTTP transport protocols, capable of processing millions of messages per second."
 edition = "2021"
 license = "Apache-2.0"
@@ -59,6 +59,7 @@ tokio = { version = "1.43.0", features = ["full"] }
 tokio-rustls = { version = "0.26.1" }
 toml = "0.8.20"
 tracing = { version = "0.1.41" }
+trait-variant = { version = "0.1.2" }
 uuid = { version = "1.13.1", features = ["v7", "fast-rng", "zerocopy"] }
 webpki-roots = { version = "0.26.8" }
 

dependencies.md

@@ -21,6 +21,7 @@ use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64};
 use std::sync::Arc;
 use std::task::{Context, Poll};
 use std::time::Duration;
+use tokio::time;
 use tokio::time::sleep;
 use tracing::{error, info, trace, warn};
 
@@ -36,10 +37,18 @@ pub enum AutoCommit {
     Disabled,
     /// The auto-commit is enabled and the offset is stored on the server after a certain interval.
     Interval(IggyDuration),
-    /// The auto-commit is enabled and the offset is stored on the server after a certain interval or depending on the mode.
+    /// The auto-commit is enabled and the offset is stored on the server after a certain interval or depending on the mode when consuming the messages.
     IntervalOrWhen(IggyDuration, AutoCommitWhen),
-    /// The auto-commit is enabled and the offset is stored on the server depending on the mode.
+    /// The auto-commit is enabled and the offset is stored on the server after a certain interval or depending on the mode after consuming the messages.
+    ///
+    /// **This will only work with the `IggyConsumerMessageExt` trait when using `consume_messages()`.**
+    IntervalOrAfter(IggyDuration, AutoCommitAfter),
+    /// The auto-commit is enabled and the offset is stored on the server depending on the mode when consuming the messages.
     When(AutoCommitWhen),
+    /// The auto-commit is enabled and the offset is stored on the server depending on the mode after consuming the messages.
+    ///
+    /// **This will only work with the `IggyConsumerMessageExt` trait when using `consume_messages()`.**
+    After(AutoCommitAfter),
 }
 
 /// The auto-commit mode for storing the offset on the server.
@@ -55,6 +64,19 @@ pub enum AutoCommitWhen {
     ConsumingEveryNthMessage(u32),
 }
 
+/// The auto-commit mode for storing the offset on the server **after** receiving the messages.
+///
+/// **This will only work with the `IggyConsumerMessageExt` trait when using `consume_messages()`.**
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum AutoCommitAfter {
+    /// The offset is stored on the server after all the messages are consumed.
+    ConsumingAllMessages,
+    /// The offset is stored on the server after consuming each message.
+    ConsumingEachMessage,
+    /// The offset is stored on the server after consuming every Nth message.
+    ConsumingEveryNthMessage(u32),
+}
+
 unsafe impl Send for IggyConsumer {}
 unsafe impl Sync for IggyConsumer {}
 
@@ -89,6 +111,8 @@ pub struct IggyConsumer {
     last_polled_at: Arc<AtomicU64>,
     current_partition_id: Arc<AtomicU32>,
     reconnection_retry_interval: IggyDuration,
+    init_retries: Option<u32>,
+    init_retry_interval: IggyDuration,
     allow_replay: bool,
 }
 
@@ -108,7 +132,9 @@ impl IggyConsumer {
         auto_join_consumer_group: bool,
         create_consumer_group_if_not_exists: bool,
         encryptor: Option<Arc<EncryptorKind>>,
-        retry_interval: IggyDuration,
+        reconnection_retry_interval: IggyDuration,
+        init_retries: Option<u32>,
+        init_retry_interval: IggyDuration,
         allow_replay: bool,
     ) -> Self {
         let (store_offset_sender, _) = flume::unbounded();
@@ -160,11 +186,17 @@ impl IggyConsumer {
             },
             last_polled_at: Arc::new(AtomicU64::new(0)),
             current_partition_id: Arc::new(AtomicU32::new(0)),
-            reconnection_retry_interval: retry_interval,
+            reconnection_retry_interval,
+            init_retries,
+            init_retry_interval,
             allow_replay,
         }
     }
 
+    pub(crate) fn auto_commit(&self) -> AutoCommit {
+        self.auto_commit
+    }
+
     /// Returns the name of the consumer.
     pub fn name(&self) -> &str {
         &self.consumer_name
@@ -230,19 +262,65 @@ impl IggyConsumer {
             return Ok(());
         }
 
+        let stream_id = self.stream_id.clone();
+        let topic_id = self.topic_id.clone();
+        let consumer_name = &self.consumer_name;
+
+        info!(
+            "Initializing consumer: {consumer_name} for stream: {stream_id}, topic: {topic_id}..."
+        );
+
         {
+            let mut retries = 0;
+            let init_retries = self.init_retries.unwrap_or_default();
+            let interval = self.init_retry_interval;
+
+            let mut timer = time::interval(interval.get_duration());
+            timer.tick().await;
+
             let client = self.client.read().await;
-            if client.get_stream(&self.stream_id).await?.is_none() {
+            let mut stream_exists = client.get_stream(&stream_id).await?.is_some();
+            let mut topic_exists = client.get_topic(&stream_id, &topic_id).await?.is_some();
+
+            loop {
+                if stream_exists && topic_exists {
+                    info!("Stream: {stream_id} and topic: {topic_id} were found. Initializing consumer...",);
+                    break;
+                }
+
+                if retries >= init_retries {
+                    break;
+                }
+
+                retries += 1;
+                if !stream_exists {
+                    warn!("Stream: {stream_id} does not exist. Retrying ({retries}/{init_retries}) in {interval}...",);
+                    timer.tick().await;
+                    stream_exists = client.get_stream(&stream_id).await?.is_some();
+                }
+
+                if !stream_exists {
+                    continue;
+                }
+
+                topic_exists = client.get_topic(&stream_id, &topic_id).await?.is_some();
+                if topic_exists {
+                    break;
+                }
+
+                warn!("Topic: {topic_id} does not exist in stream: {stream_id}. Retrying ({retries}/{init_retries}) in {interval}...",);
+                timer.tick().await;
+            }
+
+            if !stream_exists {
+                error!("Stream: {stream_id} was not found.");
                 return Err(IggyError::StreamNameNotFound(
                     self.stream_id.get_string_value().unwrap_or_default(),
                 ));
-            }
+            };
 
-            if client
-                .get_topic(&self.stream_id, &self.topic_id)
-                .await?
-                .is_none()
-            {
+            if !topic_exists {
+                error!("Topic: {topic_id} was not found in stream: {stream_id}.");
                 return Err(IggyError::TopicNameNotFound(
                     self.topic_id.get_string_value().unwrap_or_default(),
                     self.stream_id.get_string_value().unwrap_or_default(),
@@ -256,6 +334,7 @@ impl IggyConsumer {
         match self.auto_commit {
             AutoCommit::Interval(interval) => self.store_offsets_in_background(interval),
             AutoCommit::IntervalOrWhen(interval, _) => self.store_offsets_in_background(interval),
+            AutoCommit::IntervalOrAfter(interval, _) => self.store_offsets_in_background(interval),
             _ => {}
         }
 
@@ -285,6 +364,10 @@ impl IggyConsumer {
         });
 
         self.initialized = true;
+        info!(
+            "Consumer: {consumer_name} has been initialized for stream: {}, topic: {}.",
+            self.stream_id, self.topic_id
+        );
         Ok(())
     }
 
@@ -359,7 +442,7 @@ impl IggyConsumer {
         });
     }
 
-    fn send_store_offset(&self, partition_id: u32, offset: u64) {
+    pub(crate) fn send_store_offset(&self, partition_id: u32, offset: u64) {
         if let Err(error) = self.store_offset_sender.send((partition_id, offset)) {
             error!("Failed to send offset to store: {error}, please verify if `init()` on IggyConsumer object has been called.");
         }
@@ -868,6 +951,8 @@ pub struct IggyConsumerBuilder {
     create_consumer_group_if_not_exists: bool,
     encryptor: Option<Arc<EncryptorKind>>,
     reconnection_retry_interval: IggyDuration,
+    init_retries: Option<u32>,
+    init_retry_interval: IggyDuration,
     allow_replay: bool,
 }
 
@@ -901,6 +986,8 @@ impl IggyConsumerBuilder {
             encryptor,
             polling_interval,
             reconnection_retry_interval: IggyDuration::ONE_SECOND,
+            init_retries: None,
+            init_retry_interval: IggyDuration::ONE_SECOND,
             allow_replay: false,
         }
     }
@@ -941,6 +1028,13 @@ impl IggyConsumerBuilder {
         }
     }
 
+    pub fn commit_failed_messages(self) -> Self {
+        Self {
+            auto_commit: AutoCommit::Disabled,
+            ..self
+        }
+    }
+
     /// Automatically joins the consumer group if the consumer is a part of a consumer group.
     pub fn auto_join_consumer_group(self) -> Self {
         Self {
@@ -1013,6 +1107,17 @@ impl IggyConsumerBuilder {
         }
     }
 
+    /// Sets the number of retries and the interval when initializing the consumer if the stream or topic is not found.
+    /// Might be useful when the stream or topic is created dynamically by the producer.
+    /// By default, the consumer will not retry.
+    pub fn init_retries(self, retries: u32, interval: IggyDuration) -> Self {
+        Self {
+            init_retries: Some(retries),
+            init_retry_interval: interval,
+            ..self
+        }
+    }
+
     /// Allows replaying the messages, `false` by default.
     pub fn allow_replay(self) -> Self {
         Self {
@@ -1040,6 +1145,8 @@ impl IggyConsumerBuilder {
             self.create_consumer_group_if_not_exists,
             self.encryptor,
             self.reconnection_retry_interval,
+            self.init_retries,
+            self.init_retry_interval,
             self.allow_replay,
         )
     }

sdk/Cargo.toml

@@ -116,16 +116,19 @@ impl IggyProducer {
             return Ok(());
         }
 
+        let stream_id = self.stream_id.clone();
+        let topic_id = self.topic_id.clone();
+        info!("Initializing producer for stream: {stream_id} and topic: {topic_id}...");
         self.subscribe_events().await;
         let client = self.client.clone();
         let client = client.read().await;
-        if client.get_stream(&self.stream_id).await?.is_none() {
+        if client.get_stream(&stream_id).await?.is_none() {
             if !self.create_stream_if_not_exists {
                 error!("Stream does not exist and auto-creation is disabled.");
                 return Err(IggyError::StreamNameNotFound(self.stream_name.clone()));
             }
 
-            let (name, id) = match self.stream_id.kind {
+            let (name, id) = match stream_id.kind {
                 IdKind::Numeric => (
                     self.stream_name.to_owned(),
                     Some(self.stream_id.get_u32_value()?),
@@ -136,11 +139,7 @@ impl IggyProducer {
             client.create_stream(&name, id).await?;
         }
 
-        if client
-            .get_topic(&self.stream_id, &self.topic_id)
-            .await?
-            .is_none()
-        {
+        if client.get_topic(&stream_id, &topic_id).await?.is_none() {
             if !self.create_topic_if_not_exists {
                 error!("Topic does not exist and auto-creation is disabled.");
                 return Err(IggyError::TopicNameNotFound(
@@ -172,6 +171,7 @@ impl IggyProducer {
         }
 
         self.initialized = true;
+        info!("Producer has been initialized for stream: {stream_id} and topic: {topic_id}.");
         Ok(())
     }