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

s2_sdk/session/
read.rs
