aerostream_client/
consumer.rs

1use std::pin::Pin;
2use std::sync::atomic::{AtomicU64, Ordering};
3use std::sync::Arc;
4use std::task::{Context, Poll};
5use std::time::Duration;
6
7use futures_core::Stream;
8use tokio::sync::{mpsc, RwLock};
9use tracing::warn;
10
11use crate::connection::AeroConnection;
12use crate::error::AeroError;
13use crate::protocol::Record;
14
15/// Configuration options for an AeroStream consumer.
16#[derive(Debug, Clone)]
17pub struct ConsumerConfig {
18    /// Broker endpoint addresses.
19    pub bootstrap_servers: Vec<String>,
20    /// Target topic to subscribe to.
21    pub topic: String,
22    /// Target partition to consume from.
23    pub partition: u32,
24    /// Starting offset for consumption.
25    pub initial_offset: u64,
26    /// Maximum bytes per fetch request.
27    pub max_bytes: u32,
28    /// Maximum long-poll wait time on broker in milliseconds.
29    pub max_wait_ms: u32,
30    /// Optional authentication bearer token.
31    pub token: Option<String>,
32    /// Connection timeout.
33    pub connect_timeout: Duration,
34    /// Poll retry backoff duration when empty or broker throttles.
35    pub backoff_on_empty: Duration,
36}
37
38impl Default for ConsumerConfig {
39    fn default() -> Self {
40        Self {
41            bootstrap_servers: vec!["127.0.0.1:9091".to_string()],
42            topic: String::new(),
43            partition: 0,
44            initial_offset: 0,
45            max_bytes: 1024 * 1024,
46            max_wait_ms: 500,
47            token: None,
48            connect_timeout: Duration::from_secs(5),
49            backoff_on_empty: Duration::from_millis(50),
50        }
51    }
52}
53
54impl ConsumerConfig {
55    /// Create a fluent builder for ConsumerConfig.
56    pub fn builder() -> ConsumerConfigBuilder {
57        ConsumerConfigBuilder::default()
58    }
59}
60
61/// Fluent builder for constructing `ConsumerConfig` instances.
62#[derive(Default)]
63pub struct ConsumerConfigBuilder {
64    config: ConsumerConfig,
65}
66
67impl ConsumerConfigBuilder {
68    /// Set list of broker addresses.
69    pub fn bootstrap_servers(mut self, servers: Vec<String>) -> Self {
70        self.config.bootstrap_servers = servers;
71        self
72    }
73
74    /// Add a single bootstrap broker address.
75    pub fn bootstrap_server(mut self, server: impl Into<String>) -> Self {
76        self.config.bootstrap_servers = vec![server.into()];
77        self
78    }
79
80    /// Set topic to consume from.
81    pub fn topic(mut self, topic: impl Into<String>) -> Self {
82        self.config.topic = topic.into();
83        self
84    }
85
86    /// Set partition index to consume from.
87    pub fn partition(mut self, partition: u32) -> Self {
88        self.config.partition = partition;
89        self
90    }
91
92    /// Set initial starting offset.
93    pub fn initial_offset(mut self, offset: u64) -> Self {
94        self.config.initial_offset = offset;
95        self
96    }
97
98    /// Set maximum bytes to retrieve per fetch request.
99    pub fn max_bytes(mut self, max_bytes: u32) -> Self {
100        self.config.max_bytes = max_bytes;
101        self
102    }
103
104    /// Set broker long-poll wait time in milliseconds.
105    pub fn max_wait_ms(mut self, max_wait_ms: u32) -> Self {
106        self.config.max_wait_ms = max_wait_ms;
107        self
108    }
109
110    /// Configure authentication token.
111    pub fn token(mut self, token: impl Into<String>) -> Self {
112        self.config.token = Some(token.into());
113        self
114    }
115
116    /// Set connection timeout duration.
117    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
118        self.config.connect_timeout = timeout;
119        self
120    }
121
122    /// Validate and build the `ConsumerConfig`.
123    pub fn build(self) -> Result<ConsumerConfig, AeroError> {
124        if self.config.bootstrap_servers.is_empty() {
125            return Err(AeroError::InvalidConfiguration(
126                "bootstrap_servers cannot be empty".into(),
127            ));
128        }
129        if self.config.topic.is_empty() {
130            return Err(AeroError::InvalidConfiguration("topic cannot be empty".into()));
131        }
132        Ok(self.config)
133    }
134}
135
136/// Continuous stream of records yielding `Result<Record, AeroError>`.
137pub struct RecordStream {
138    receiver: mpsc::Receiver<Result<Record, AeroError>>,
139    worker_handle: tokio::task::JoinHandle<()>,
140}
141
142impl Drop for RecordStream {
143    fn drop(&mut self) {
144        self.worker_handle.abort();
145    }
146}
147
148impl Stream for RecordStream {
149    type Item = Result<Record, AeroError>;
150
151    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
152        self.receiver.poll_recv(cx)
153    }
154}
155
156/// Asynchronous consumer client for fetching messages from AeroStream.
157#[derive(Clone)]
158pub struct AeroConsumer {
159    config: Arc<ConsumerConfig>,
160    connection: Arc<RwLock<Option<AeroConnection>>>,
161    current_offset: Arc<AtomicU64>,
162}
163
164impl AeroConsumer {
165    /// Create a new `AeroConsumer` with the specified configuration.
166    pub fn new(config: ConsumerConfig) -> Self {
167        let initial_offset = config.initial_offset;
168        Self {
169            config: Arc::new(config),
170            connection: Arc::new(RwLock::new(None)),
171            current_offset: Arc::new(AtomicU64::new(initial_offset)),
172        }
173    }
174
175    /// Get current consumption offset.
176    pub fn offset(&self) -> u64 {
177        self.current_offset.load(Ordering::Acquire)
178    }
179
180    /// Set consumption offset.
181    pub fn seek(&self, offset: u64) {
182        self.current_offset.store(offset, Ordering::Release);
183    }
184
185    /// Obtain an active connection or establish a new one.
186    async fn get_or_connect(&self) -> Result<AeroConnection, AeroError> {
187        {
188            let read = self.connection.read().await;
189            if let Some(conn) = &*read {
190                if conn.is_alive() {
191                    return Ok(conn.clone());
192                }
193            }
194        }
195
196        let mut write = self.connection.write().await;
197        if let Some(conn) = &*write {
198            if conn.is_alive() {
199                return Ok(conn.clone());
200            }
201        }
202
203        let mut last_err = AeroError::Disconnected;
204        for server in &self.config.bootstrap_servers {
205            match AeroConnection::connect(
206                server,
207                self.config.token.as_deref(),
208                self.config.connect_timeout,
209            )
210            .await
211            {
212                Ok(conn) => {
213                    *write = Some(conn.clone());
214                    return Ok(conn);
215                }
216                Err(e) => {
217                    warn!(broker = %server, error = %e, "Failed connecting to broker");
218                    last_err = e;
219                }
220            }
221        }
222
223        Err(last_err)
224    }
225
226    /// Fetch a batch of records using Command 4 (Multi-Entry Long-Polling Fetch).
227    pub async fn fetch_multi(
228        &self,
229        topic: &str,
230        partition: u32,
231        start_offset: u64,
232        max_bytes: u32,
233        max_wait_ms: u32,
234    ) -> Result<Vec<Record>, AeroError> {
235        let conn = self.get_or_connect().await?;
236        match conn
237            .fetch_multi(topic, partition, start_offset, max_bytes, max_wait_ms)
238            .await
239        {
240            Ok(records) => Ok(records),
241            Err(e) => {
242                conn.mark_broken();
243                Err(e)
244            }
245        }
246    }
247
248    /// Poll for next available records and automatically advance internal offset.
249    pub async fn poll_next_batch(&self) -> Result<Vec<Record>, AeroError> {
250        let offset = self.offset();
251        let records = self
252            .fetch_multi(
253                &self.config.topic,
254                self.config.partition,
255                offset,
256                self.config.max_bytes,
257                self.config.max_wait_ms,
258            )
259            .await?;
260
261        if let Some(last) = records.last() {
262            self.seek(last.offset + 1);
263        }
264        Ok(records)
265    }
266
267    /// Create a continuous asynchronous stream of records.
268    pub fn stream(&self) -> RecordStream {
269        let (tx, rx) = mpsc::channel(256);
270        let consumer = self.clone();
271
272        let worker_handle = tokio::spawn(async move {
273            loop {
274                match consumer.poll_next_batch().await {
275                    Ok(records) => {
276                        if records.is_empty() {
277                            tokio::time::sleep(consumer.config.backoff_on_empty).await;
278                            continue;
279                        }
280
281                        for record in records {
282                            if tx.send(Ok(record)).await.is_err() {
283                                return; // Downstream receiver dropped
284                            }
285                        }
286                    }
287                    Err(err) => {
288                        let is_fatal = matches!(&err, AeroError::AuthFailed(_));
289                        if tx.send(Err(err)).await.is_err() || is_fatal {
290                            return;
291                        }
292                        tokio::time::sleep(Duration::from_millis(200)).await;
293                    }
294                }
295            }
296        });
297
298        RecordStream {
299            receiver: rx,
300            worker_handle,
301        }
302    }
303}
aerostream_client/consumer.rs

aerostream_client/
consumer.rs
