laser_sdk/
cursor.rs

1use crate::error::LaserError;
2use crate::laser::Laser;
3use crate::message::Message;
4use crate::types::MessageId;
5use iggy::prelude::*;
6use std::collections::BTreeMap;
7
8const DEFAULT_BATCH: u32 = 1000;
9
10impl Laser {
11    /// A resumable, offset-addressable reader over `topic`. Each [`poll`](Cursor::poll)
12    /// drains every partition from where the previous poll stopped, so repeated
13    /// calls read only what is new, never a rescan from zero. You own the offsets:
14    /// read them with [`offsets`](Cursor::offsets), persist them in any
15    /// [`StateStore`](crate::state_store::StateStore) (in-memory, file, or managed
16    /// `Kv`), and resume after a restart with [`from_offsets`](Cursor::from_offsets).
17    ///
18    /// This is the open streaming primitive the `Agent` runtime sits above: reach
19    /// for the runtime when you want Apache Iggy to track offsets for you (consumer
20    /// groups, dedup, DLQ), and for a `Cursor` when you want to own the cursor and
21    /// checkpoint it yourself.
22    ///
23    /// Reads from the default stream. Use [`reader_on`](Self::reader_on) to read a
24    /// topic on an explicit stream.
25    #[cfg(feature = "agent")]
26    pub(crate) fn reader(&self, topic: impl Into<String>) -> Result<Cursor, LaserError> {
27        let stream = self.stream_required()?.to_owned();
28        Cursor::new(self, &stream, &topic.into())
29    }
30
31    /// A resumable reader over `topic` on an explicit `stream`, for connection-only
32    /// handles or reading across several streams from one connection.
33    pub(crate) fn reader_on(
34        &self,
35        stream: impl Into<String>,
36        topic: impl Into<String>,
37    ) -> Result<Cursor, LaserError> {
38        Cursor::new(self, &stream.into(), &topic.into())
39    }
40}
41
42/// A resumable reader over one topic. Build it with [`Topic::replay`]. Owns a
43/// cheap `Laser` clone, so it moves into a `tokio::spawn` or is stored beside the
44/// `Laser` rather than borrowing it.
45///
46/// [`Topic::replay`]: crate::stream::Topic::replay
47pub struct Cursor {
48    laser: Laser,
49    stream: Identifier,
50    topic: Identifier,
51    consumer: Consumer,
52    offsets: Vec<u64>,
53    batch: u32,
54}
55
56impl Cursor {
57    fn new(laser: &Laser, stream: &str, topic: &str) -> Result<Self, LaserError> {
58        Ok(Self {
59            stream: Identifier::named(stream)?,
60            topic: Identifier::named(topic)?,
61            consumer: Consumer::new(Identifier::named("laser-cursor")?),
62            laser: laser.clone(),
63            offsets: Vec::new(),
64            batch: DEFAULT_BATCH,
65        })
66    }
67
68    /// Resume from previously persisted per-partition offsets: the next offset to
69    /// read on each partition, exactly what an earlier [`offsets`](Self::offsets)
70    /// returned. A shorter vec than the topic's partition count is padded with 0
71    /// (read those partitions from the start).
72    pub fn from_offsets(mut self, offsets: Vec<u64>) -> Self {
73        self.offsets = offsets;
74        self
75    }
76
77    /// Size each server poll at `batch` messages per partition (default 1000).
78    ///
79    /// This is the per-request size, not a ceiling on one [`poll`](Self::poll):
80    /// a poll keeps requesting until it reaches the partition's tail or the
81    /// per-partition drain ceiling, whichever comes first. Persist
82    /// [`offsets`](Self::offsets) and poll again to continue past the ceiling.
83    pub fn batch(mut self, batch: u32) -> Self {
84        self.batch = batch.max(1);
85        self
86    }
87
88    // Attribute the reads to a caller-chosen reader identity (the typed
89    // reader passes its name through here). Offsets stay client-owned
90    // either way, the identity only names the reader server-side.
91    pub(crate) fn consumer_named(mut self, name: &str) -> Result<Self, LaserError> {
92        self.consumer = Consumer::new(Identifier::named(name)?);
93        Ok(self)
94    }
95
96    /// The next offset to read on each partition. Persist this to resume later with
97    /// [`from_offsets`](Self::from_offsets).
98    pub fn offsets(&self) -> &[u64] {
99        &self.offsets
100    }
101
102    /// This cursor as a [`Stream`](futures::Stream) of messages, one at a time,
103    /// draining everything currently appended and ending (like the async reader in
104    /// the Python binding) once caught up. A later `stream()` on a fresh cursor
105    /// seeded from the persisted [`offsets`](Self::offsets) resumes from there. Use
106    /// it with `futures::StreamExt`: `while let Some(message) = stream.next().await`.
107    /// The bounded-read primitive as a stream, not a substitute for the reliable
108    /// consumer's continuous delivery.
109    pub fn stream(self) -> impl futures::Stream<Item = Result<Message, LaserError>> {
110        futures::stream::unfold(
111            (self, std::collections::VecDeque::new(), false),
112            |(mut cursor, mut buffered, stop)| async move {
113                if stop {
114                    return None;
115                }
116                loop {
117                    if let Some(message) = buffered.pop_front() {
118                        return Some((Ok(message), (cursor, buffered, false)));
119                    }
120                    match cursor.poll().await {
121                        Ok(batch) if batch.is_empty() => return None,
122                        Ok(batch) => buffered.extend(batch),
123                        Err(error) => return Some((Err(error), (cursor, buffered, true))),
124                    }
125                }
126            },
127        )
128    }
129
130    /// Drain everything appended since the last poll, advancing the cursor. Returns
131    /// the new messages ordered by Iggy timestamp (then partition, offset), or an
132    /// empty vec when the reader is caught up (or the topic does not exist yet).
133    #[tracing::instrument(target = "laser", level = "debug", skip_all, fields(topic = %self.topic, operation = "poll"))]
134    pub async fn poll(&mut self) -> Result<Vec<Message>, LaserError> {
135        Ok(self
136            .poll_iggy()
137            .await?
138            .into_iter()
139            .map(|(partition, message)| Message {
140                id: MessageId::new(partition, message.header.offset),
141                payload: message.payload.to_vec(),
142                headers: headers_to_strings(&message),
143            })
144            .collect())
145    }
146
147    pub(crate) async fn poll_iggy(&mut self) -> Result<Vec<(u32, IggyMessage)>, LaserError> {
148        let Some(details) = self
149            .laser
150            .client()
151            .get_topic(&self.stream, &self.topic)
152            .await?
153        else {
154            return Ok(Vec::new());
155        };
156        let partitions = crate::poll::bounded_partitions(details.partitions_count) as usize;
157        let mut next_offsets = self.offsets.clone();
158        if next_offsets.len() < partitions {
159            next_offsets.resize(partitions, 0);
160        }
161        // (timestamp, partition, offset, message) so the merge across partitions is
162        // ordered by Iggy's single clock, like `ContextAssembler`.
163        let mut collected: Vec<(u64, u32, u64, IggyMessage)> = Vec::new();
164        for partition in 0..partitions as u32 {
165            let batch = crate::poll::drain_partition(
166                &self.laser.client(),
167                &self.stream,
168                &self.topic,
169                &self.consumer,
170                crate::poll::DrainRange::open(partition, next_offsets[partition as usize]),
171                self.batch,
172            )
173            .await?;
174            next_offsets[partition as usize] = batch.next_offset;
175            for message in batch.messages {
176                let offset = message.header.offset;
177                collected.push((message.header.timestamp, partition, offset, message));
178            }
179        }
180        collected
181            .sort_by_key(|(timestamp, partition, offset, _)| (*timestamp, *partition, *offset));
182        self.offsets = next_offsets;
183        Ok(collected
184            .into_iter()
185            .map(|(_, partition, _, message)| (partition, message))
186            .collect())
187    }
188}
189
190// The message's user headers decoded to strings. Non-UTF-8 keys/values are
191// dropped (the agent layer reconstructs `Provenance` from these same headers).
192fn headers_to_strings(message: &IggyMessage) -> BTreeMap<String, String> {
193    let Ok(Some(headers)) = message.user_headers_map() else {
194        return BTreeMap::new();
195    };
196    headers
197        .iter()
198        .filter_map(|(key, value)| {
199            Some((
200                key.as_str().ok()?.to_owned(),
201                value.as_str().ok()?.to_owned(),
202            ))
203        })
204        .collect()
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210    use axum::extract::{Query, State};
211    use axum::http::StatusCode;
212    use axum::response::{IntoResponse, Response};
213    use axum::routing::get;
214    use axum::{Json, Router};
215    use iggy::http::http_client::HttpClient;
216    use std::collections::HashMap;
217    use std::sync::Arc;
218    use std::sync::atomic::{AtomicBool, Ordering};
219    use tokio::net::TcpListener;
220    use tokio::sync::Notify;
221    use tokio::task::JoinHandle;
222
223    #[tokio::test]
224    async fn given_later_partition_failure_when_polling_then_should_preserve_all_offsets() {
225        let server = TestServer::start(2, 2).await;
226        server.state.fail.store(true, Ordering::SeqCst);
227        let mut cursor = server.cursor().from_offsets(vec![0, 0]);
228
229        assert!(cursor.poll().await.is_err());
230        assert_eq!(cursor.offsets(), &[0, 0]);
231
232        let messages = cursor.poll().await.expect("retry all partitions");
233        assert_eq!(messages.len(), 4);
234        assert_eq!(cursor.offsets(), &[2, 2]);
235    }
236
237    #[tokio::test]
238    async fn given_later_partition_wait_when_poll_is_cancelled_then_should_preserve_all_offsets() {
239        let server = TestServer::start(2, 2).await;
240        server.state.block.store(true, Ordering::SeqCst);
241        let mut cursor = server.cursor().from_offsets(vec![0, 0]);
242
243        tokio::select! {
244            result = cursor.poll() => panic!("poll must wait for the second partition: {result:?}"),
245            () = server.state.waiting.notified() => {}
246        }
247        assert_eq!(cursor.offsets(), &[0, 0]);
248
249        server.state.block.store(false, Ordering::SeqCst);
250        server.state.release.notify_one();
251        let messages = cursor.poll().await.expect("retry cancelled poll");
252        assert_eq!(messages.len(), 4);
253        assert_eq!(cursor.offsets(), &[2, 2]);
254    }
255
256    #[tokio::test]
257    async fn given_large_batches_when_polling_then_should_bound_reads_and_resume_without_gaps() {
258        let limit = crate::poll::MAX_DRAIN_MESSAGES;
259        let server = TestServer::start(1, limit as u64 + 1).await;
260        for batch in [6000, 20000] {
261            let mut cursor = server.cursor().batch(batch);
262            let messages = cursor.poll().await.expect("bounded read");
263            assert_eq!(messages.len(), limit, "batch size {batch}");
264            assert_eq!(cursor.offsets(), &[limit as u64]);
265
266            let remaining = cursor.poll().await.expect("resume after limit");
267            assert_eq!(remaining.len(), 1);
268            assert_eq!(remaining[0].id, MessageId::new(0, limit as u64));
269            assert_eq!(cursor.offsets(), &[limit as u64 + 1]);
270        }
271    }
272
273    struct TestServer {
274        url: String,
275        state: Arc<TestState>,
276        task: JoinHandle<()>,
277    }
278
279    struct TestState {
280        partitions: u32,
281        messages: u64,
282        fail: AtomicBool,
283        block: AtomicBool,
284        waiting: Notify,
285        release: Notify,
286    }
287
288    impl TestServer {
289        async fn start(partitions: u32, messages: u64) -> Self {
290            let state = Arc::new(TestState {
291                partitions,
292                messages,
293                fail: AtomicBool::new(false),
294                block: AtomicBool::new(false),
295                waiting: Notify::new(),
296                release: Notify::new(),
297            });
298            let app = Router::new()
299                .route("/streams/test/topics/events", get(topic))
300                .route("/streams/test/topics/events/messages", get(poll))
301                .with_state(state.clone());
302            let listener = TcpListener::bind("127.0.0.1:0")
303                .await
304                .expect("test listener");
305            let url = format!(
306                "http://{}",
307                listener.local_addr().expect("listener address")
308            );
309            let task = tokio::spawn(async move {
310                axum::serve(listener, app).await.expect("test server");
311            });
312            Self { url, state, task }
313        }
314
315        fn cursor(&self) -> Cursor {
316            let client = HttpClient::create(Arc::new(HttpClientConfig {
317                api_url: self.url.clone(),
318                jwt: Some("test-token".to_owned()),
319                retries: 0,
320                ..Default::default()
321            }))
322            .expect("HTTP client");
323            Laser::from_client(IggyClient::new(ClientWrapper::Http(client)))
324                .stream("test")
325                .topic("events")
326                .replay()
327                .expect("cursor")
328        }
329    }
330
331    impl Drop for TestServer {
332        fn drop(&mut self) {
333            self.task.abort();
334        }
335    }
336
337    async fn topic(State(state): State<Arc<TestState>>) -> Json<TopicDetails> {
338        Json(TopicDetails {
339            id: 1,
340            created_at: IggyTimestamp::from(0),
341            name: "events".to_owned(),
342            size: IggyByteSize::from(0),
343            message_expiry: IggyExpiry::NeverExpire,
344            compression_algorithm: CompressionAlgorithm::None,
345            max_topic_size: MaxTopicSize::ServerDefault,
346            messages_count: state.messages * u64::from(state.partitions),
347            partitions_count: state.partitions,
348            partitions: Vec::new(),
349            options: Default::default(),
350        })
351    }
352
353    async fn poll(
354        State(state): State<Arc<TestState>>,
355        Query(query): Query<HashMap<String, String>>,
356    ) -> Response {
357        let partition = query["partition_id"].parse().expect("partition id");
358        if partition == 1 {
359            if state.fail.swap(false, Ordering::SeqCst) {
360                return (StatusCode::BAD_REQUEST, "invalid poll").into_response();
361            }
362            if state.block.load(Ordering::SeqCst) {
363                state.waiting.notify_one();
364                state.release.notified().await;
365            }
366        }
367        let offset: u64 = query["value"].parse().expect("poll offset");
368        let count: u64 = query["count"].parse().expect("poll count");
369        let messages: Vec<_> = (offset..(offset + count).min(state.messages))
370            .map(|offset| {
371                let mut message = IggyMessage::builder()
372                    .payload(vec![1].into())
373                    .build()
374                    .expect("test message");
375                message.header.offset = offset;
376                message.header.timestamp = offset;
377                message
378            })
379            .collect();
380        Json(PolledMessages {
381            partition_id: partition,
382            current_offset: state.messages.saturating_sub(1),
383            count: messages.len() as u32,
384            messages,
385        })
386        .into_response()
387    }
388}
laser_sdk/cursor.rs

laser_sdk/
cursor.rs
