laser-sdk 0.6.0

use crate::error::LaserError;
use crate::laser::Laser;
use crate::message::Message;
use crate::types::MessageId;
use iggy::prelude::*;
use std::collections::BTreeMap;

const DEFAULT_BATCH: u32 = 1000;

impl Laser {
    /// A resumable, offset-addressable reader over `topic`. Each [`poll`](Cursor::poll)
    /// drains every partition from where the previous poll stopped, so repeated
    /// calls read only what is new, never a rescan from zero. You own the offsets:
    /// read them with [`offsets`](Cursor::offsets), persist them in any
    /// [`StateStore`](crate::state_store::StateStore) (in-memory, file, or managed
    /// `Kv`), and resume after a restart with [`from_offsets`](Cursor::from_offsets).
    ///
    /// This is the open streaming primitive the `Agent` runtime sits above: reach
    /// for the runtime when you want Apache Iggy to track offsets for you (consumer
    /// groups, dedup, DLQ), and for a `Cursor` when you want to own the cursor and
    /// checkpoint it yourself.
    ///
    /// Reads from the default stream. Use [`reader_on`](Self::reader_on) to read a
    /// topic on an explicit stream.
    #[cfg(feature = "agent")]
    pub(crate) fn reader(&self, topic: impl Into<String>) -> Result<Cursor, LaserError> {
        let stream = self.stream_required()?.to_owned();
        Cursor::new(self, &stream, &topic.into())
    }

    /// A resumable reader over `topic` on an explicit `stream`, for connection-only
    /// handles or reading across several streams from one connection.
    pub(crate) fn reader_on(
        &self,
        stream: impl Into<String>,
        topic: impl Into<String>,
    ) -> Result<Cursor, LaserError> {
        Cursor::new(self, &stream.into(), &topic.into())
    }
}

/// A resumable reader over one topic. Build it with [`Topic::replay`]. Owns a
/// cheap `Laser` clone, so it moves into a `tokio::spawn` or is stored beside the
/// `Laser` rather than borrowing it.
///
/// [`Topic::replay`]: crate::stream::Topic::replay
pub struct Cursor {
    laser: Laser,
    stream: Identifier,
    topic: Identifier,
    consumer: Consumer,
    offsets: Vec<u64>,
    batch: u32,
}

impl Cursor {
    fn new(laser: &Laser, stream: &str, topic: &str) -> Result<Self, LaserError> {
        Ok(Self {
            stream: Identifier::named(stream)?,
            topic: Identifier::named(topic)?,
            consumer: Consumer::new(Identifier::named("laser-cursor")?),
            laser: laser.clone(),
            offsets: Vec::new(),
            batch: DEFAULT_BATCH,
        })
    }

    /// Resume from previously persisted per-partition offsets: the next offset to
    /// read on each partition, exactly what an earlier [`offsets`](Self::offsets)
    /// returned. A shorter vec than the topic's partition count is padded with 0
    /// (read those partitions from the start).
    pub fn from_offsets(mut self, offsets: Vec<u64>) -> Self {
        self.offsets = offsets;
        self
    }

    /// Size each server poll at `batch` messages per partition (default 1000).
    ///
    /// This is the per-request size, not a ceiling on one [`poll`](Self::poll):
    /// a poll keeps requesting until it reaches the partition's tail or the
    /// per-partition drain ceiling, whichever comes first. Persist
    /// [`offsets`](Self::offsets) and poll again to continue past the ceiling.
    pub fn batch(mut self, batch: u32) -> Self {
        self.batch = batch.max(1);
        self
    }

    // Attribute the reads to a caller-chosen reader identity (the typed
    // reader passes its name through here). Offsets stay client-owned
    // either way, the identity only names the reader server-side.
    pub(crate) fn consumer_named(mut self, name: &str) -> Result<Self, LaserError> {
        self.consumer = Consumer::new(Identifier::named(name)?);
        Ok(self)
    }

    /// The next offset to read on each partition. Persist this to resume later with
    /// [`from_offsets`](Self::from_offsets).
    pub fn offsets(&self) -> &[u64] {
        &self.offsets
    }

    /// This cursor as a [`Stream`](futures::Stream) of messages, one at a time,
    /// draining everything currently appended and ending (like the async reader in
    /// the Python binding) once caught up. A later `stream()` on a fresh cursor
    /// seeded from the persisted [`offsets`](Self::offsets) resumes from there. Use
    /// it with `futures::StreamExt`: `while let Some(message) = stream.next().await`.
    /// The bounded-read primitive as a stream, not a substitute for the reliable
    /// consumer's continuous delivery.
    pub fn stream(self) -> impl futures::Stream<Item = Result<Message, LaserError>> {
        futures::stream::unfold(
            (self, std::collections::VecDeque::new(), false),
            |(mut cursor, mut buffered, stop)| async move {
                if stop {
                    return None;
                }
                loop {
                    if let Some(message) = buffered.pop_front() {
                        return Some((Ok(message), (cursor, buffered, false)));
                    }
                    match cursor.poll().await {
                        Ok(batch) if batch.is_empty() => return None,
                        Ok(batch) => buffered.extend(batch),
                        Err(error) => return Some((Err(error), (cursor, buffered, true))),
                    }
                }
            },
        )
    }

    /// Drain everything appended since the last poll, advancing the cursor. Returns
    /// the new messages ordered by Iggy timestamp (then partition, offset), or an
    /// empty vec when the reader is caught up (or the topic does not exist yet).
    #[tracing::instrument(target = "laser", level = "debug", skip_all, fields(topic = %self.topic, operation = "poll"))]
    pub async fn poll(&mut self) -> Result<Vec<Message>, LaserError> {
        Ok(self
            .poll_iggy()
            .await?
            .into_iter()
            .map(|(partition, message)| Message {
                id: MessageId::new(partition, message.header.offset),
                payload: message.payload.to_vec(),
                headers: headers_to_strings(&message),
            })
            .collect())
    }

    pub(crate) async fn poll_iggy(&mut self) -> Result<Vec<(u32, IggyMessage)>, LaserError> {
        let Some(details) = self
            .laser
            .client()
            .get_topic(&self.stream, &self.topic)
            .await?
        else {
            return Ok(Vec::new());
        };
        let partitions = crate::poll::bounded_partitions(details.partitions_count) as usize;
        let mut next_offsets = self.offsets.clone();
        if next_offsets.len() < partitions {
            next_offsets.resize(partitions, 0);
        }
        // (timestamp, partition, offset, message) so the merge across partitions is
        // ordered by Iggy's single clock, like `ContextAssembler`.
        let mut collected: Vec<(u64, u32, u64, IggyMessage)> = Vec::new();
        for partition in 0..partitions as u32 {
            let batch = crate::poll::drain_partition(
                &self.laser.client(),
                &self.stream,
                &self.topic,
                &self.consumer,
                crate::poll::DrainRange::open(partition, next_offsets[partition as usize]),
                self.batch,
            )
            .await?;
            next_offsets[partition as usize] = batch.next_offset;
            for message in batch.messages {
                let offset = message.header.offset;
                collected.push((message.header.timestamp, partition, offset, message));
            }
        }
        collected
            .sort_by_key(|(timestamp, partition, offset, _)| (*timestamp, *partition, *offset));
        self.offsets = next_offsets;
        Ok(collected
            .into_iter()
            .map(|(_, partition, _, message)| (partition, message))
            .collect())
    }
}

// The message's user headers decoded to strings. Non-UTF-8 keys/values are
// dropped (the agent layer reconstructs `Provenance` from these same headers).
fn headers_to_strings(message: &IggyMessage) -> BTreeMap<String, String> {
    let Ok(Some(headers)) = message.user_headers_map() else {
        return BTreeMap::new();
    };
    headers
        .iter()
        .filter_map(|(key, value)| {
            Some((
                key.as_str().ok()?.to_owned(),
                value.as_str().ok()?.to_owned(),
            ))
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::extract::{Query, State};
    use axum::http::StatusCode;
    use axum::response::{IntoResponse, Response};
    use axum::routing::get;
    use axum::{Json, Router};
    use iggy::http::http_client::HttpClient;
    use std::collections::HashMap;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicBool, Ordering};
    use tokio::net::TcpListener;
    use tokio::sync::Notify;
    use tokio::task::JoinHandle;

    #[tokio::test]
    async fn given_later_partition_failure_when_polling_then_should_preserve_all_offsets() {
        let server = TestServer::start(2, 2).await;
        server.state.fail.store(true, Ordering::SeqCst);
        let mut cursor = server.cursor().from_offsets(vec![0, 0]);

        assert!(cursor.poll().await.is_err());
        assert_eq!(cursor.offsets(), &[0, 0]);

        let messages = cursor.poll().await.expect("retry all partitions");
        assert_eq!(messages.len(), 4);
        assert_eq!(cursor.offsets(), &[2, 2]);
    }

    #[tokio::test]
    async fn given_later_partition_wait_when_poll_is_cancelled_then_should_preserve_all_offsets() {
        let server = TestServer::start(2, 2).await;
        server.state.block.store(true, Ordering::SeqCst);
        let mut cursor = server.cursor().from_offsets(vec![0, 0]);

        tokio::select! {
            result = cursor.poll() => panic!("poll must wait for the second partition: {result:?}"),
            () = server.state.waiting.notified() => {}
        }
        assert_eq!(cursor.offsets(), &[0, 0]);

        server.state.block.store(false, Ordering::SeqCst);
        server.state.release.notify_one();
        let messages = cursor.poll().await.expect("retry cancelled poll");
        assert_eq!(messages.len(), 4);
        assert_eq!(cursor.offsets(), &[2, 2]);
    }

    #[tokio::test]
    async fn given_large_batches_when_polling_then_should_bound_reads_and_resume_without_gaps() {
        let limit = crate::poll::MAX_DRAIN_MESSAGES;
        let server = TestServer::start(1, limit as u64 + 1).await;
        for batch in [6000, 20000] {
            let mut cursor = server.cursor().batch(batch);
            let messages = cursor.poll().await.expect("bounded read");
            assert_eq!(messages.len(), limit, "batch size {batch}");
            assert_eq!(cursor.offsets(), &[limit as u64]);

            let remaining = cursor.poll().await.expect("resume after limit");
            assert_eq!(remaining.len(), 1);
            assert_eq!(remaining[0].id, MessageId::new(0, limit as u64));
            assert_eq!(cursor.offsets(), &[limit as u64 + 1]);
        }
    }

    struct TestServer {
        url: String,
        state: Arc<TestState>,
        task: JoinHandle<()>,
    }

    struct TestState {
        partitions: u32,
        messages: u64,
        fail: AtomicBool,
        block: AtomicBool,
        waiting: Notify,
        release: Notify,
    }

    impl TestServer {
        async fn start(partitions: u32, messages: u64) -> Self {
            let state = Arc::new(TestState {
                partitions,
                messages,
                fail: AtomicBool::new(false),
                block: AtomicBool::new(false),
                waiting: Notify::new(),
                release: Notify::new(),
            });
            let app = Router::new()
                .route("/streams/test/topics/events", get(topic))
                .route("/streams/test/topics/events/messages", get(poll))
                .with_state(state.clone());
            let listener = TcpListener::bind("127.0.0.1:0")
                .await
                .expect("test listener");
            let url = format!(
                "http://{}",
                listener.local_addr().expect("listener address")
            );
            let task = tokio::spawn(async move {
                axum::serve(listener, app).await.expect("test server");
            });
            Self { url, state, task }
        }

        fn cursor(&self) -> Cursor {
            let client = HttpClient::create(Arc::new(HttpClientConfig {
                api_url: self.url.clone(),
                jwt: Some("test-token".to_owned()),
                retries: 0,
                ..Default::default()
            }))
            .expect("HTTP client");
            Laser::from_client(IggyClient::new(ClientWrapper::Http(client)))
                .stream("test")
                .topic("events")
                .replay()
                .expect("cursor")
        }
    }

    impl Drop for TestServer {
        fn drop(&mut self) {
            self.task.abort();
        }
    }

    async fn topic(State(state): State<Arc<TestState>>) -> Json<TopicDetails> {
        Json(TopicDetails {
            id: 1,
            created_at: IggyTimestamp::from(0),
            name: "events".to_owned(),
            size: IggyByteSize::from(0),
            message_expiry: IggyExpiry::NeverExpire,
            compression_algorithm: CompressionAlgorithm::None,
            max_topic_size: MaxTopicSize::ServerDefault,
            messages_count: state.messages * u64::from(state.partitions),
            partitions_count: state.partitions,
            partitions: Vec::new(),
            options: Default::default(),
        })
    }

    async fn poll(
        State(state): State<Arc<TestState>>,
        Query(query): Query<HashMap<String, String>>,
    ) -> Response {
        let partition = query["partition_id"].parse().expect("partition id");
        if partition == 1 {
            if state.fail.swap(false, Ordering::SeqCst) {
                return (StatusCode::BAD_REQUEST, "invalid poll").into_response();
            }
            if state.block.load(Ordering::SeqCst) {
                state.waiting.notify_one();
                state.release.notified().await;
            }
        }
        let offset: u64 = query["value"].parse().expect("poll offset");
        let count: u64 = query["count"].parse().expect("poll count");
        let messages: Vec<_> = (offset..(offset + count).min(state.messages))
            .map(|offset| {
                let mut message = IggyMessage::builder()
                    .payload(vec![1].into())
                    .build()
                    .expect("test message");
                message.header.offset = offset;
                message.header.timestamp = offset;
                message
            })
            .collect();
        Json(PolledMessages {
            partition_id: partition,
            current_offset: state.messages.saturating_sub(1),
            count: messages.len() as u32,
            messages,
        })
        .into_response()
    }
}