s2_sdk/session/
read.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    task::{Context, Poll},
5    time::Duration,
6};
7
8use async_stream::{stream, try_stream};
9use futures_util::{
10    StreamExt,
11    future::{FutureExt, Shared},
12};
13use s2_api::v1::stream::{ReadEnd, ReadStart};
14use tokio::{
15    sync::oneshot,
16    time::{Instant, timeout},
17};
18use tracing::debug;
19
20use crate::{
21    api::{ApiError, BasinClient, retry_builder},
22    retry::RetryBackoff,
23    types::{EncryptionKey, MeteredBytes, ReadBatch, S2Error, StreamName, StreamPosition},
24};
25
26#[derive(Debug, thiserror::Error)]
27pub enum ReadSessionError {
28    #[error(transparent)]
29    Api(#[from] ApiError),
30    #[error("heartbeat timeout")]
31    HeartbeatTimeout,
32}
33
34impl ReadSessionError {
35    pub fn is_retryable(&self) -> bool {
36        match self {
37            Self::Api(err) => err.is_retryable(),
38            Self::HeartbeatTimeout => true,
39        }
40    }
41}
42
43impl From<ReadSessionError> for S2Error {
44    fn from(err: ReadSessionError) -> Self {
45        match err {
46            ReadSessionError::Api(api_err) => api_err.into(),
47            other => S2Error::Client(other.to_string()),
48        }
49    }
50}
51
52pub type Streaming<R> =
53    Pin<Box<dyn Send + futures_core::Stream<Item = Result<R, ReadSessionError>>>>;
54
55#[derive(Debug, Clone, thiserror::Error)]
56#[non_exhaustive]
57/// Error returned while waiting for a read session to catch up.
58pub enum CaughtUpError {
59    #[error("read session ended before catching up")]
60    /// The session ended before reaching a reported tail.
61    SessionClosed,
62    #[error(transparent)]
63    /// The read failed.
64    Read(#[from] S2Error),
65}
66
67impl From<CaughtUpError> for S2Error {
68    fn from(err: CaughtUpError) -> Self {
69        match err {
70            CaughtUpError::SessionClosed => {
71                Self::Client("read session ended before catching up".into())
72            }
73            CaughtUpError::Read(err) => err,
74        }
75    }
76}
77
78type CaughtUpResult = Result<StreamPosition, CaughtUpError>;
79
80#[derive(Clone)]
81enum CaughtUpFuture {
82    Pending(Shared<oneshot::Receiver<CaughtUpResult>>),
83    Ready(CaughtUpResult),
84}
85
86impl Future for CaughtUpFuture {
87    type Output = CaughtUpResult;
88
89    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
90        match &mut *self {
91            Self::Pending(future) => match Pin::new(future).poll(cx) {
92                Poll::Ready(Ok(result)) => Poll::Ready(result),
93                Poll::Ready(Err(_)) => Poll::Ready(Err(CaughtUpError::SessionClosed)),
94                Poll::Pending => Poll::Pending,
95            },
96            Self::Ready(result) => Poll::Ready(result.clone()),
97        }
98    }
99}
100
101struct CaughtUpState {
102    /// Latest reported tail we've fully delivered, if currently caught up.
103    tail: Option<StreamPosition>,
104    /// Once set, the session has ended.
105    terminal: bool,
106    /// Fires the current caught-up future.
107    tx: Option<oneshot::Sender<CaughtUpResult>>,
108    /// The future handed out by `caught_up()`.
109    future: CaughtUpFuture,
110}
111
112impl CaughtUpState {
113    fn new() -> Self {
114        let (tx, future) = pending_catch_up();
115        Self {
116            tail: None,
117            terminal: false,
118            tx: Some(tx),
119            future,
120        }
121    }
122
123    fn is_caught_up(&self) -> bool {
124        self.tail.is_some()
125    }
126
127    fn future(&self) -> CaughtUpFuture {
128        self.future.clone()
129    }
130
131    fn set_behind(&mut self) {
132        if self.terminal || self.tail.take().is_none() {
133            return;
134        }
135        let (tx, future) = pending_catch_up();
136        self.tx = Some(tx);
137        self.future = future;
138    }
139
140    fn set_caught_up(&mut self, tail: StreamPosition) {
141        if self.terminal || self.tail == Some(tail) {
142            return;
143        }
144        self.tail = Some(tail);
145        self.complete(Ok(tail));
146    }
147
148    fn end(&mut self, error: Option<S2Error>) {
149        if self.terminal {
150            return;
151        }
152        self.terminal = true;
153        if let Some(error) = error {
154            self.tail = None;
155            self.complete(Err(CaughtUpError::Read(error)));
156        } else if self.tail.is_none() {
157            self.complete(Err(CaughtUpError::SessionClosed));
158        }
159    }
160
161    fn complete(&mut self, result: CaughtUpResult) {
162        if let Some(tx) = self.tx.take() {
163            let _ = tx.send(result);
164        } else {
165            self.future = CaughtUpFuture::Ready(result);
166        }
167    }
168}
169
170fn pending_catch_up() -> (oneshot::Sender<CaughtUpResult>, CaughtUpFuture) {
171    let (tx, rx) = oneshot::channel();
172    (tx, CaughtUpFuture::Pending(rx.shared()))
173}
174
175struct ReadUpdate {
176    batch: Option<ReadBatch>,
177    caught_up_tail: Option<StreamPosition>,
178}
179
180impl ReadUpdate {
181    fn behind() -> Self {
182        Self {
183            batch: None,
184            caught_up_tail: None,
185        }
186    }
187
188    fn from_batch(mut batch: ReadBatch, ignore_command_records: bool) -> Self {
189        let caught_up_tail = batch.tail.filter(|tail| {
190            batch.records.is_empty()
191                || batch
192                    .records
193                    .last()
194                    .is_some_and(|record| record.seq_num.checked_add(1) == Some(tail.seq_num))
195        });
196
197        if ignore_command_records {
198            batch.records.retain(|record| !record.is_command_record());
199        }
200
201        Self {
202            batch: (!batch.records.is_empty()).then_some(batch),
203            caught_up_tail,
204        }
205    }
206}
207
208/// A continuous stream of read batches.
209pub struct ReadSession {
210    updates: Streaming<ReadUpdate>,
211    state: CaughtUpState,
212}
213
214impl ReadSession {
215    fn new(updates: Streaming<ReadUpdate>) -> Self {
216        Self {
217            updates,
218            state: CaughtUpState::new(),
219        }
220    }
221
222    /// Return whether all records through the latest reported tail were delivered.
223    ///
224    /// A later batch that does not reach a reported tail or a reconnect resets it.
225    /// Ignored command records count toward progress. Use
226    /// [`S2Stream::check_tail`](crate::S2Stream::check_tail) for the current tail.
227    pub fn is_caught_up(&self) -> bool {
228        self.state.is_caught_up()
229    }
230
231    /// Return a future for the current or next caught-up tail.
232    ///
233    /// Continue polling the read session while awaiting this future; the future does not drive
234    /// reads itself. It is ready immediately when the session is already caught up and remains
235    /// pending across retries. Once it resolves, its returned tail never changes. If the session
236    /// later falls behind, call `caught_up()` again to wait for the next catch-up. The future
237    /// returns [`CaughtUpError`] if the session fails or closes before catching up.
238    pub fn caught_up(
239        &self,
240    ) -> impl Future<Output = Result<StreamPosition, CaughtUpError>>
241    + Clone
242    + Send
243    + Sync
244    + Unpin
245    + 'static {
246        self.state.future()
247    }
248}
249
250impl futures_core::Stream for ReadSession {
251    type Item = Result<ReadBatch, S2Error>;
252
253    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
254        loop {
255            match self.updates.as_mut().poll_next(cx) {
256                Poll::Pending => return Poll::Pending,
257                Poll::Ready(Some(Ok(update))) => {
258                    if let Some(tail) = update.caught_up_tail {
259                        self.state.set_caught_up(tail);
260                    } else {
261                        self.state.set_behind();
262                    }
263                    if let Some(batch) = update.batch {
264                        return Poll::Ready(Some(Ok(batch)));
265                    }
266                }
267                Poll::Ready(Some(Err(error))) => {
268                    let error = S2Error::from(error);
269                    self.state.end(Some(error.clone()));
270                    return Poll::Ready(Some(Err(error)));
271                }
272                Poll::Ready(None) => {
273                    self.state.end(None);
274                    return Poll::Ready(None);
275                }
276            }
277        }
278    }
279}
280
281impl Drop for ReadSession {
282    fn drop(&mut self) {
283        self.state.end(None);
284    }
285}
286
287pub async fn read_session(
288    client: BasinClient,
289    name: StreamName,
290    encryption: Option<EncryptionKey>,
291    mut start: ReadStart,
292    mut end: ReadEnd,
293    ignore_command_records: bool,
294) -> Result<ReadSession, ReadSessionError> {
295    let mut retry_backoff = retry_builder(&client.config.retry).build();
296    let baseline_wait = end.wait;
297    let mut last_tail_at: Option<Instant> = None;
298
299    let batches = loop {
300        end.wait = remaining_wait(baseline_wait, last_tail_at);
301        match session_inner(
302            client.clone(),
303            name.clone(),
304            encryption.clone(),
305            start.clone(),
306            end.clone(),
307        )
308        .await
309        {
310            Ok(batches) => {
311                retry_backoff.reset();
312                break batches;
313            }
314            Err(err) => {
315                if let Some(backoff) = retry_delay(&err, &mut retry_backoff) {
316                    tokio::time::sleep(backoff).await;
317                    continue;
318                }
319                return Err(err);
320            }
321        }
322    };
323
324    let updates = Box::pin(stream! {
325        let mut batches: Option<Streaming<ReadBatch>> = Some(batches);
326
327        loop {
328            if batches.is_none() {
329                end.wait = remaining_wait(baseline_wait, last_tail_at);
330                match session_inner(
331                    client.clone(),
332                    name.clone(),
333                    encryption.clone(),
334                    start.clone(),
335                    end.clone(),
336                ).await {
337                    Ok(b) => batches = Some(b),
338                    Err(err) => {
339                        if let Some(backoff) = retry_delay(&err, &mut retry_backoff) {
340                            tokio::time::sleep(backoff).await;
341                            continue;
342                        }
343                        yield Err(err);
344                        break;
345                    }
346                }
347            }
348
349            match batches
350                .as_mut()
351                .expect("batches should not be None")
352                .next()
353                .await
354            {
355                Some(Ok(batch)) => {
356                    if retry_backoff.used() > 0 {
357                        retry_backoff.reset();
358                    }
359
360                    if batch.tail.is_some() {
361                        last_tail_at = Some(Instant::now());
362                    }
363
364                    if let Some(record) = batch.records.last() {
365                        start = ReadStart {
366                            seq_num: Some(record.seq_num + 1),
367                            timestamp: None,
368                            tail_offset: None,
369                            clamp: start.clamp,
370                        };
371                    }
372                    if let Some(count) = end.count.as_mut() {
373                        *count = count.saturating_sub(batch.records.len())
374                    }
375                    if let Some(bytes) = end.bytes.as_mut() {
376                        *bytes = bytes.saturating_sub(
377                            batch.records.iter().map(|r| r.metered_bytes()).sum()
378                        )
379                    }
380
381                    yield Ok(ReadUpdate::from_batch(batch, ignore_command_records));
382                }
383                Some(Err(err)) => {
384                    batches = None;
385                    if let Some(backoff) = retry_delay(&err, &mut retry_backoff) {
386                        yield Ok(ReadUpdate::behind());
387                        tokio::time::sleep(backoff).await;
388                        continue;
389                    }
390                    yield Err(err);
391                    break;
392                }
393                None => break,
394            }
395        }
396    });
397    Ok(ReadSession::new(updates))
398}
399
400async fn session_inner(
401    client: BasinClient,
402    name: StreamName,
403    encryption: Option<EncryptionKey>,
404    start: ReadStart,
405    end: ReadEnd,
406) -> Result<Streaming<ReadBatch>, ReadSessionError> {
407    let mut batches = client
408        .read_session(&name, start, end, encryption.as_ref())
409        .await?;
410    Ok(Box::pin(try_stream! {
411        loop {
412            match timeout(Duration::from_secs(20), batches.next()).await {
413                Ok(Some(batch)) => {
414                    yield ReadBatch::from_api(batch?);
415                }
416                Ok(None) => break,
417                Err(_) => Err(ReadSessionError::HeartbeatTimeout)?,
418            }
419        }
420    }))
421}
422
423/// Compute the remaining wait budget for a retry.
424///
425/// During catchup (tail not yet observed), the full wait is sent.
426/// Once tailing, the wait budget is depleted based on time since
427/// the last batch with tail info, which approximates how long the
428/// server has been in its long polling state.
429fn remaining_wait(baseline_wait: Option<u32>, last_tail_at: Option<Instant>) -> Option<u32> {
430    baseline_wait.map(|w| match last_tail_at {
431        Some(since) => w.saturating_sub(since.elapsed().as_secs() as u32),
432        None => w,
433    })
434}
435
436fn retry_delay(err: &ReadSessionError, backoffs: &mut RetryBackoff) -> Option<Duration> {
437    if err.is_retryable()
438        && let Some(backoff) = backoffs.next()
439    {
440        debug!(
441            %err,
442            ?backoff,
443            num_retries_remaining = backoffs.remaining(),
444            "retrying read session"
445        );
446        Some(backoff)
447    } else {
448        debug!(
449            %err,
450            is_retryable = err.is_retryable(),
451            retries_exhausted = backoffs.is_exhausted(),
452            "not retrying read session"
453        );
454        None
455    }
456}
457
458#[cfg(test)]
459mod tests {
460    use bytes::Bytes;
461    use futures_util::{StreamExt, poll, stream};
462    use tokio::sync::mpsc;
463    use tokio_stream::wrappers::UnboundedReceiverStream;
464
465    use super::*;
466    use crate::types::{Header, SequencedRecord};
467
468    fn position(seq_num: u64) -> StreamPosition {
469        StreamPosition {
470            seq_num,
471            timestamp: seq_num,
472        }
473    }
474
475    fn record(seq_num: u64, command: bool) -> SequencedRecord {
476        SequencedRecord {
477            seq_num,
478            timestamp: seq_num,
479            body: Bytes::new(),
480            headers: if command {
481                vec![Header::new("", "fence")]
482            } else {
483                Vec::new()
484            },
485        }
486    }
487
488    fn batch(records: Vec<SequencedRecord>, tail: Option<StreamPosition>) -> ReadBatch {
489        ReadBatch { records, tail }
490    }
491
492    fn test_session(
493        updates: impl futures_core::Stream<Item = Result<ReadUpdate, ReadSessionError>> + Send + 'static,
494    ) -> ReadSession {
495        ReadSession::new(Box::pin(updates))
496    }
497
498    #[tokio::test]
499    async fn caught_up_follows_delivery_and_pins_tail() {
500        let tail = position(2);
501        let mut session = test_session(stream::iter([
502            Ok(ReadUpdate::from_batch(
503                batch(vec![record(0, false), record(1, false)], Some(tail)),
504                false,
505            )),
506            Ok(ReadUpdate::from_batch(
507                batch(vec![record(2, false)], Some(position(5))),
508                false,
509            )),
510        ]));
511        let caught_up = session.caught_up();
512        let mut pending = Box::pin(caught_up.clone());
513
514        assert!(poll!(pending.as_mut()).is_pending());
515        assert!(!session.is_caught_up());
516
517        let first = session.next().await.unwrap().unwrap();
518        assert_eq!(first.records.len(), 2);
519        assert!(session.is_caught_up());
520        let caught_up_while_caught = session.caught_up();
521
522        session.next().await.unwrap().unwrap();
523        assert!(!session.is_caught_up());
524        assert_eq!(caught_up.await.unwrap(), tail);
525        assert_eq!(caught_up_while_caught.await.unwrap(), tail);
526    }
527
528    #[tokio::test]
529    async fn heartbeat_waits_for_visible_batch() {
530        let tail = position(2);
531        let (tx, rx) = mpsc::unbounded_channel();
532        let mut session = test_session(UnboundedReceiverStream::new(rx));
533        let caught_up = session.caught_up();
534
535        tx.send(Ok(ReadUpdate::from_batch(
536            batch(vec![record(0, false), record(1, false)], None),
537            false,
538        )))
539        .unwrap();
540        tx.send(Ok(ReadUpdate::from_batch(
541            batch(Vec::new(), Some(tail)),
542            false,
543        )))
544        .unwrap();
545
546        assert_eq!(session.next().await.unwrap().unwrap().records.len(), 2);
547        assert!(!session.is_caught_up());
548
549        let mut next = Box::pin(session.next());
550        assert!(poll!(next.as_mut()).is_pending());
551        drop(next);
552        assert!(session.is_caught_up());
553        assert_eq!(caught_up.await.unwrap(), tail);
554    }
555
556    #[tokio::test]
557    async fn unchanged_heartbeat_reuses_caught_up_future() {
558        let tail = position(1);
559        let (tx, rx) = mpsc::unbounded_channel();
560        let mut session = test_session(UnboundedReceiverStream::new(rx));
561
562        tx.send(Ok(ReadUpdate::from_batch(
563            batch(vec![record(0, false)], Some(tail)),
564            false,
565        )))
566        .unwrap();
567        session.next().await.unwrap().unwrap();
568        let caught_up = session.state.future();
569
570        tx.send(Ok(ReadUpdate::from_batch(
571            batch(Vec::new(), Some(tail)),
572            false,
573        )))
574        .unwrap();
575        let mut next = Box::pin(session.next());
576        assert!(poll!(next.as_mut()).is_pending());
577        drop(next);
578
579        let CaughtUpFuture::Pending(caught_up) = caught_up else {
580            panic!("initial caught-up future should use the pending epoch");
581        };
582        let CaughtUpFuture::Pending(current) = session.state.future() else {
583            panic!("unchanged heartbeat should preserve the pending epoch");
584        };
585        assert!(caught_up.ptr_eq(&current));
586    }
587
588    #[tokio::test]
589    async fn filtered_command_counts_toward_caught_up() {
590        let tail = position(2);
591        let mut session = test_session(stream::iter([
592            Ok(ReadUpdate::from_batch(
593                batch(vec![record(0, false)], None),
594                true,
595            )),
596            Ok(ReadUpdate::from_batch(
597                batch(vec![record(1, true)], Some(tail)),
598                true,
599            )),
600        ]));
601        let caught_up = session.caught_up();
602
603        let delivered = session.next().await.unwrap().unwrap();
604        assert_eq!(delivered.records.len(), 1);
605        assert_eq!(delivered.records[0].seq_num, 0);
606        assert!(!session.is_caught_up());
607
608        assert!(session.next().await.is_none());
609        assert!(session.is_caught_up());
610        assert_eq!(caught_up.await.unwrap(), tail);
611    }
612
613    #[tokio::test]
614    async fn caught_up_wait_survives_retry() {
615        let first_tail = position(1);
616        let tail = position(3);
617        let (tx, rx) = mpsc::unbounded_channel();
618        let mut session = test_session(UnboundedReceiverStream::new(rx));
619
620        tx.send(Ok(ReadUpdate::from_batch(
621            batch(Vec::new(), Some(first_tail)),
622            false,
623        )))
624        .unwrap();
625        let mut next = Box::pin(session.next());
626        assert!(poll!(next.as_mut()).is_pending());
627        drop(next);
628        assert!(session.is_caught_up());
629
630        tx.send(Ok(ReadUpdate::behind())).unwrap();
631        let mut next = Box::pin(session.next());
632        assert!(poll!(next.as_mut()).is_pending());
633        drop(next);
634        assert!(!session.is_caught_up());
635        let caught_up = session.caught_up();
636
637        tx.send(Ok(ReadUpdate::behind())).unwrap();
638        tx.send(Ok(ReadUpdate::from_batch(
639            batch(Vec::new(), Some(tail)),
640            false,
641        )))
642        .unwrap();
643        drop(tx);
644        assert!(session.next().await.is_none());
645        assert_eq!(caught_up.await.unwrap(), tail);
646    }
647
648    #[tokio::test]
649    async fn clean_end_rejects_wait() {
650        let mut session = test_session(stream::empty());
651        let caught_up = session.caught_up();
652
653        assert!(session.next().await.is_none());
654        assert!(matches!(caught_up.await, Err(CaughtUpError::SessionClosed)));
655    }
656
657    #[tokio::test]
658    async fn read_error_rejects_wait() {
659        let mut session = test_session(stream::iter([Err(ReadSessionError::HeartbeatTimeout)]));
660        let caught_up = session.caught_up();
661
662        let error = session.next().await.unwrap().unwrap_err();
663        assert_eq!(error.to_string(), "heartbeat timeout");
664        assert!(matches!(
665            caught_up.await,
666            Err(CaughtUpError::Read(S2Error::Client(message)))
667                if message == "heartbeat timeout"
668        ));
669    }
670
671    #[tokio::test]
672    async fn read_error_after_caught_up_preserves_resolved_future() {
673        let tail = position(1);
674        let mut session = test_session(stream::iter([
675            Ok(ReadUpdate::from_batch(
676                batch(vec![record(0, false)], Some(tail)),
677                false,
678            )),
679            Err(ReadSessionError::HeartbeatTimeout),
680        ]));
681
682        session.next().await.unwrap().unwrap();
683        assert!(session.is_caught_up());
684        let caught_up = session.caught_up();
685
686        session.next().await.unwrap().unwrap_err();
687        assert!(!session.is_caught_up());
688        assert_eq!(caught_up.await.unwrap(), tail);
689        assert!(matches!(
690            session.caught_up().await,
691            Err(CaughtUpError::Read(S2Error::Client(message)))
692                if message == "heartbeat timeout"
693        ));
694    }
695
696    #[tokio::test]
697    async fn dropping_session_rejects_wait() {
698        let caught_up = {
699            let session = test_session(stream::pending());
700            session.caught_up()
701        };
702
703        assert!(matches!(caught_up.await, Err(CaughtUpError::SessionClosed)));
704    }
705}
s2_sdk/session/read.rs

s2_sdk/session/
read.rs
