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    resume_seq_num: Option<u64>,
220}
221
222impl ReadUpdate {
223    fn behind() -> Self {
224        Self {
225            batch: None,
226            caught_up_tail: None,
227            resume_seq_num: None,
228        }
229    }
230
231    fn from_batch(mut batch: ReadBatch, ignore_command_records: bool) -> Self {
232        let resume_seq_num = resume_seq_num_after_batch(&batch);
233        let caught_up_tail = batch.tail.filter(|tail| {
234            batch.records.is_empty()
235                || batch
236                    .records
237                    .last()
238                    .is_some_and(|record| record.seq_num.checked_add(1) == Some(tail.seq_num))
239        });
240
241        if ignore_command_records {
242            batch.records.retain(|record| !record.is_command_record());
243        }
244
245        Self {
246            batch: (!batch.records.is_empty()).then_some(batch),
247            caught_up_tail,
248            resume_seq_num,
249        }
250    }
251}
252
253/// A continuous stream of read batches.
254pub struct ReadSession {
255    updates: InternalStreaming<ReadUpdate>,
256    state: CaughtUpState,
257    resume_seq_num: Option<u64>,
258}
259
260impl ReadSession {
261    fn new(updates: InternalStreaming<ReadUpdate>, resume_seq_num: Option<u64>) -> Self {
262        Self {
263            updates,
264            state: CaughtUpState::new(),
265            resume_seq_num,
266        }
267    }
268
269    /// Return the absolute sequence number from which the session would resume after a retry.
270    ///
271    /// An unclamped absolute starting sequence number is available immediately. A timestamp,
272    /// tail-relative, or clamped start returns `None` until the session receives a record or a
273    /// reported tail. The returned value is the sequence number of the next record the session
274    /// expects. It advances as the session is polled, including across records hidden by
275    /// [`ReadInput::ignore_command_records`](crate::types::ReadInput::ignore_command_records).
276    pub fn resume_seq_num(&self) -> Option<u64> {
277        self.resume_seq_num
278    }
279
280    /// Return whether all records through the latest reported tail were delivered.
281    ///
282    /// A later batch that does not reach a reported tail or a reconnect resets it.
283    /// Ignored command records count toward progress. Use
284    /// [`S2Stream::check_tail`](crate::S2Stream::check_tail) for the current tail.
285    pub fn is_caught_up(&self) -> bool {
286        self.state.is_caught_up()
287    }
288
289    /// Return a future for the current or next caught-up tail.
290    ///
291    /// Continue polling the read session while awaiting this future; the future does not drive
292    /// reads itself. It is ready immediately when the session is already caught up and remains
293    /// pending across retries. Once it resolves, its returned tail never changes. If the session
294    /// later falls behind, call `caught_up()` again to wait for the next catch-up. The future
295    /// returns [`CaughtUpError`] if the session fails or closes before catching up.
296    pub fn caught_up(
297        &self,
298    ) -> impl Future<Output = Result<StreamPosition, CaughtUpError>>
299    + Clone
300    + Send
301    + Sync
302    + Unpin
303    + 'static {
304        self.state.future()
305    }
306}
307
308impl futures_core::Stream for ReadSession {
309    type Item = Result<ReadBatch, ReadSessionError>;
310
311    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
312        loop {
313            match self.updates.as_mut().poll_next(cx) {
314                Poll::Pending => return Poll::Pending,
315                Poll::Ready(Some(Ok(update))) => {
316                    if let Some(resume_seq_num) = update.resume_seq_num {
317                        self.resume_seq_num = Some(resume_seq_num);
318                    }
319                    if let Some(tail) = update.caught_up_tail {
320                        self.state.set_caught_up(tail);
321                    } else {
322                        self.state.set_behind();
323                    }
324                    if let Some(batch) = update.batch {
325                        return Poll::Ready(Some(Ok(batch)));
326                    }
327                }
328                Poll::Ready(Some(Err(error))) => {
329                    let error = ReadSessionError::from(error);
330                    self.state.end(Some(error.clone()));
331                    return Poll::Ready(Some(Err(error)));
332                }
333                Poll::Ready(None) => {
334                    self.state.end(None);
335                    return Poll::Ready(None);
336                }
337            }
338        }
339    }
340}
341
342impl Drop for ReadSession {
343    fn drop(&mut self) {
344        self.state.end(None);
345    }
346}
347
348pub async fn read_session(
349    client: BasinClient,
350    name: StreamName,
351    encryption: Option<EncryptionKey>,
352    input: ReadInput,
353    config: ReadSessionConfig,
354) -> Result<ReadSession, ReadSessionError> {
355    let ReadInput {
356        start,
357        stop,
358        ignore_command_records,
359    } = input;
360    let mut start: ReadStart = start.into();
361    let mut end: ReadEnd = stop.into();
362    let retry_policy = config.retry_policy;
363    let mut retry_backoff = retry_builder(&client.config.retry).build();
364    let baseline_wait = end.wait;
365    let mut last_tail_at: Option<Instant> = None;
366    let initial_resume_seq_num = if start.clamp == Some(true) {
367        None
368    } else {
369        start.seq_num
370    };
371
372    let batches = loop {
373        end.wait = remaining_wait(baseline_wait, last_tail_at);
374        match session_inner(
375            client.clone(),
376            name.clone(),
377            encryption.clone(),
378            start.clone(),
379            end.clone(),
380        )
381        .await
382        {
383            Ok(batches) => {
384                retry_backoff.reset();
385                break batches;
386            }
387            Err(err) => {
388                if let Some(backoff) = retry_delay(&err, &mut retry_backoff, retry_policy) {
389                    tokio::time::sleep(backoff).await;
390                    continue;
391                }
392                return Err(err.into());
393            }
394        }
395    };
396
397    let updates = Box::pin(stream! {
398        let mut batches: Option<InternalStreaming<ReadBatch>> = Some(batches);
399
400        loop {
401            if batches.is_none() {
402                end.wait = remaining_wait(baseline_wait, last_tail_at);
403                match session_inner(
404                    client.clone(),
405                    name.clone(),
406                    encryption.clone(),
407                    start.clone(),
408                    end.clone(),
409                ).await {
410                    Ok(b) => batches = Some(b),
411                    Err(err) => {
412                        if let Some(backoff) =
413                            retry_delay(&err, &mut retry_backoff, retry_policy)
414                        {
415                            tokio::time::sleep(backoff).await;
416                            continue;
417                        }
418                        yield Err(err);
419                        break;
420                    }
421                }
422            }
423
424            match batches
425                .as_mut()
426                .expect("batches should not be None")
427                .next()
428                .await
429            {
430                Some(Ok(batch)) => {
431                    if retry_backoff.used() > 0 {
432                        retry_backoff.reset();
433                    }
434
435                    if batch.tail.is_some() {
436                        last_tail_at = Some(Instant::now());
437                    }
438
439                    update_resume_start(&mut start, &batch);
440                    if let Some(count) = end.count.as_mut() {
441                        *count = count.saturating_sub(batch.records.len())
442                    }
443                    if let Some(bytes) = end.bytes.as_mut() {
444                        *bytes = bytes.saturating_sub(
445                            batch.records.iter().map(|r| r.metered_bytes()).sum()
446                        )
447                    }
448
449                    yield Ok(ReadUpdate::from_batch(batch, ignore_command_records));
450                }
451                Some(Err(err)) => {
452                    batches = None;
453                    if let Some(backoff) =
454                        retry_delay(&err, &mut retry_backoff, retry_policy)
455                    {
456                        yield Ok(ReadUpdate::behind());
457                        tokio::time::sleep(backoff).await;
458                        continue;
459                    }
460                    yield Err(err);
461                    break;
462                }
463                None => break,
464            }
465        }
466    });
467    Ok(ReadSession::new(updates, initial_resume_seq_num))
468}
469
470fn resume_seq_num_after_batch(batch: &ReadBatch) -> Option<u64> {
471    batch
472        .records
473        .last()
474        .map(|record| record.seq_num + 1)
475        .or_else(|| batch.tail.as_ref().map(|tail| tail.seq_num))
476}
477
478/// Advance the absolute start used when reconnecting the read session.
479///
480/// An empty batch with a reported tail still resolves a relative or timestamp start. Anchoring it
481/// prevents a reconnect from evaluating the original start against a newer tail.
482fn update_resume_start(start: &mut ReadStart, batch: &ReadBatch) {
483    if let Some(seq_num) = resume_seq_num_after_batch(batch) {
484        *start = ReadStart {
485            seq_num: Some(seq_num),
486            timestamp: None,
487            tail_offset: None,
488            clamp: start.clamp,
489        };
490    }
491}
492
493async fn session_inner(
494    client: BasinClient,
495    name: StreamName,
496    encryption: Option<EncryptionKey>,
497    start: ReadStart,
498    end: ReadEnd,
499) -> Result<InternalStreaming<ReadBatch>, ReadSessionFailure> {
500    let mut batches = client
501        .read_session(&name, start, end, encryption.as_ref())
502        .await?;
503    Ok(Box::pin(try_stream! {
504        loop {
505            match timeout(Duration::from_secs(20), batches.next()).await {
506                Ok(Some(batch)) => {
507                    yield ReadBatch::from_api(batch?);
508                }
509                Ok(None) => break,
510                Err(_) => Err(ReadSessionFailure::HeartbeatTimeout)?,
511            }
512        }
513    }))
514}
515
516/// Compute the remaining wait budget for a retry.
517///
518/// During catchup (tail not yet observed), the full wait is sent.
519/// Once tailing, the wait budget is depleted based on time since
520/// the last batch with tail info, which approximates how long the
521/// server has been in its long polling state.
522fn remaining_wait(baseline_wait: Option<u32>, last_tail_at: Option<Instant>) -> Option<u32> {
523    baseline_wait.map(|w| match last_tail_at {
524        Some(since) => w.saturating_sub(since.elapsed().as_secs() as u32),
525        None => w,
526    })
527}
528
529fn retry_delay(
530    err: &ReadSessionFailure,
531    backoffs: &mut RetryBackoff,
532    retry_policy: ReadSessionRetryPolicy,
533) -> Option<Duration> {
534    if !err.is_retryable() {
535        debug!(
536            %err,
537            is_retryable = false,
538            retries_exhausted = backoffs.is_exhausted(),
539            "not retrying read session"
540        );
541        return None;
542    }
543
544    let backoff = match retry_policy {
545        ReadSessionRetryPolicy::Budgeted => backoffs.next(),
546        ReadSessionRetryPolicy::Indefinite => Some(backoffs.next_or_max()),
547    };
548    if let Some(backoff) = backoff {
549        debug!(
550            %err,
551            ?backoff,
552            ?retry_policy,
553            num_retries_remaining = backoffs.remaining(),
554            "retrying read session"
555        );
556        Some(backoff)
557    } else {
558        debug!(
559            %err,
560            is_retryable = err.is_retryable(),
561            retries_exhausted = backoffs.is_exhausted(),
562            "not retrying read session"
563        );
564        None
565    }
566}
567
568#[cfg(test)]
569mod tests {
570    use bytes::Bytes;
571    use futures_util::{StreamExt, poll, stream};
572    use tokio::sync::mpsc;
573    use tokio_stream::wrappers::UnboundedReceiverStream;
574
575    use super::*;
576    use crate::types::{Header, SequencedRecord};
577
578    fn position(seq_num: u64) -> StreamPosition {
579        StreamPosition {
580            seq_num,
581            timestamp: seq_num,
582        }
583    }
584
585    fn record(seq_num: u64, command: bool) -> SequencedRecord {
586        SequencedRecord {
587            seq_num,
588            timestamp: seq_num,
589            body: Bytes::new(),
590            headers: if command {
591                vec![Header::new("", "fence")]
592            } else {
593                Vec::new()
594            },
595        }
596    }
597
598    fn batch(records: Vec<SequencedRecord>, tail: Option<StreamPosition>) -> ReadBatch {
599        ReadBatch { records, tail }
600    }
601
602    #[test]
603    fn empty_tail_anchors_relative_resume_start() {
604        let mut start = ReadStart {
605            seq_num: None,
606            timestamp: None,
607            tail_offset: Some(0),
608            clamp: Some(true),
609        };
610
611        update_resume_start(&mut start, &batch(Vec::new(), Some(position(42))));
612
613        assert_eq!(start.seq_num, Some(42));
614        assert_eq!(start.timestamp, None);
615        assert_eq!(start.tail_offset, None);
616        assert_eq!(start.clamp, Some(true));
617    }
618
619    fn test_session(
620        updates: impl futures_core::Stream<Item = Result<ReadUpdate, ReadSessionFailure>>
621        + Send
622        + 'static,
623    ) -> ReadSession {
624        ReadSession::new(Box::pin(updates), None)
625    }
626
627    #[tokio::test]
628    async fn empty_tail_exposes_absolute_resume_seq_num() {
629        let (tx, rx) = mpsc::unbounded_channel();
630        let mut session = test_session(UnboundedReceiverStream::new(rx));
631
632        assert_eq!(session.resume_seq_num(), None);
633        tx.send(Ok(ReadUpdate::from_batch(
634            batch(Vec::new(), Some(position(42))),
635            false,
636        )))
637        .unwrap();
638
639        let mut next = Box::pin(session.next());
640        assert!(poll!(next.as_mut()).is_pending());
641        drop(next);
642
643        assert_eq!(session.resume_seq_num(), Some(42));
644    }
645
646    #[tokio::test]
647    async fn caught_up_follows_delivery_and_pins_tail() {
648        let tail = position(2);
649        let mut session = test_session(stream::iter([
650            Ok(ReadUpdate::from_batch(
651                batch(vec![record(0, false), record(1, false)], Some(tail)),
652                false,
653            )),
654            Ok(ReadUpdate::from_batch(
655                batch(vec![record(2, false)], Some(position(5))),
656                false,
657            )),
658        ]));
659        let caught_up = session.caught_up();
660        let mut pending = Box::pin(caught_up.clone());
661
662        assert!(poll!(pending.as_mut()).is_pending());
663        assert!(!session.is_caught_up());
664
665        let first = session.next().await.unwrap().unwrap();
666        assert_eq!(first.records.len(), 2);
667        assert!(session.is_caught_up());
668        assert_eq!(session.resume_seq_num(), Some(2));
669        let caught_up_while_caught = session.caught_up();
670
671        session.next().await.unwrap().unwrap();
672        assert!(!session.is_caught_up());
673        assert_eq!(session.resume_seq_num(), Some(3));
674        assert_eq!(caught_up.await.unwrap(), tail);
675        assert_eq!(caught_up_while_caught.await.unwrap(), tail);
676    }
677
678    #[tokio::test]
679    async fn heartbeat_waits_for_visible_batch() {
680        let tail = position(2);
681        let (tx, rx) = mpsc::unbounded_channel();
682        let mut session = test_session(UnboundedReceiverStream::new(rx));
683        let caught_up = session.caught_up();
684
685        tx.send(Ok(ReadUpdate::from_batch(
686            batch(vec![record(0, false), record(1, false)], None),
687            false,
688        )))
689        .unwrap();
690        tx.send(Ok(ReadUpdate::from_batch(
691            batch(Vec::new(), Some(tail)),
692            false,
693        )))
694        .unwrap();
695
696        assert_eq!(session.next().await.unwrap().unwrap().records.len(), 2);
697        assert!(!session.is_caught_up());
698
699        let mut next = Box::pin(session.next());
700        assert!(poll!(next.as_mut()).is_pending());
701        drop(next);
702        assert!(session.is_caught_up());
703        assert_eq!(caught_up.await.unwrap(), tail);
704    }
705
706    #[tokio::test]
707    async fn unchanged_heartbeat_reuses_caught_up_future() {
708        let tail = position(1);
709        let (tx, rx) = mpsc::unbounded_channel();
710        let mut session = test_session(UnboundedReceiverStream::new(rx));
711
712        tx.send(Ok(ReadUpdate::from_batch(
713            batch(vec![record(0, false)], Some(tail)),
714            false,
715        )))
716        .unwrap();
717        session.next().await.unwrap().unwrap();
718        let caught_up = session.state.future();
719
720        tx.send(Ok(ReadUpdate::from_batch(
721            batch(Vec::new(), Some(tail)),
722            false,
723        )))
724        .unwrap();
725        let mut next = Box::pin(session.next());
726        assert!(poll!(next.as_mut()).is_pending());
727        drop(next);
728
729        let CaughtUpFuture::Pending(caught_up) = caught_up else {
730            panic!("initial caught-up future should use the pending epoch");
731        };
732        let CaughtUpFuture::Pending(current) = session.state.future() else {
733            panic!("unchanged heartbeat should preserve the pending epoch");
734        };
735        assert!(caught_up.ptr_eq(&current));
736    }
737
738    #[tokio::test]
739    async fn filtered_command_counts_toward_caught_up() {
740        let tail = position(2);
741        let mut session = test_session(stream::iter([
742            Ok(ReadUpdate::from_batch(
743                batch(vec![record(0, false)], None),
744                true,
745            )),
746            Ok(ReadUpdate::from_batch(
747                batch(vec![record(1, true)], Some(tail)),
748                true,
749            )),
750        ]));
751        let caught_up = session.caught_up();
752
753        let delivered = session.next().await.unwrap().unwrap();
754        assert_eq!(delivered.records.len(), 1);
755        assert_eq!(delivered.records[0].seq_num, 0);
756        assert!(!session.is_caught_up());
757
758        assert!(session.next().await.is_none());
759        assert!(session.is_caught_up());
760        assert_eq!(session.resume_seq_num(), Some(2));
761        assert_eq!(caught_up.await.unwrap(), tail);
762    }
763
764    #[tokio::test]
765    async fn caught_up_wait_survives_retry() {
766        let first_tail = position(1);
767        let tail = position(3);
768        let (tx, rx) = mpsc::unbounded_channel();
769        let mut session = test_session(UnboundedReceiverStream::new(rx));
770
771        tx.send(Ok(ReadUpdate::from_batch(
772            batch(Vec::new(), Some(first_tail)),
773            false,
774        )))
775        .unwrap();
776        let mut next = Box::pin(session.next());
777        assert!(poll!(next.as_mut()).is_pending());
778        drop(next);
779        assert!(session.is_caught_up());
780
781        tx.send(Ok(ReadUpdate::behind())).unwrap();
782        let mut next = Box::pin(session.next());
783        assert!(poll!(next.as_mut()).is_pending());
784        drop(next);
785        assert!(!session.is_caught_up());
786        let caught_up = session.caught_up();
787
788        tx.send(Ok(ReadUpdate::behind())).unwrap();
789        tx.send(Ok(ReadUpdate::from_batch(
790            batch(Vec::new(), Some(tail)),
791            false,
792        )))
793        .unwrap();
794        drop(tx);
795        assert!(session.next().await.is_none());
796        assert_eq!(caught_up.await.unwrap(), tail);
797    }
798
799    #[tokio::test]
800    async fn clean_end_rejects_wait() {
801        let mut session = test_session(stream::empty());
802        let caught_up = session.caught_up();
803
804        assert!(session.next().await.is_none());
805        assert!(matches!(caught_up.await, Err(CaughtUpError::SessionClosed)));
806    }
807
808    #[tokio::test]
809    async fn read_error_rejects_wait() {
810        let mut session = test_session(stream::iter([Err(ReadSessionFailure::HeartbeatTimeout)]));
811        let caught_up = session.caught_up();
812
813        let error = session.next().await.unwrap().unwrap_err();
814        assert_eq!(error.to_string(), "heartbeat timeout");
815        assert!(matches!(
816            caught_up.await,
817            Err(CaughtUpError::Read(ReadSessionError::HeartbeatTimeout))
818        ));
819    }
820
821    #[tokio::test]
822    async fn read_error_after_caught_up_preserves_resolved_future() {
823        let tail = position(1);
824        let mut session = test_session(stream::iter([
825            Ok(ReadUpdate::from_batch(
826                batch(vec![record(0, false)], Some(tail)),
827                false,
828            )),
829            Err(ReadSessionFailure::HeartbeatTimeout),
830        ]));
831
832        session.next().await.unwrap().unwrap();
833        assert!(session.is_caught_up());
834        let caught_up = session.caught_up();
835
836        session.next().await.unwrap().unwrap_err();
837        assert!(!session.is_caught_up());
838        assert_eq!(caught_up.await.unwrap(), tail);
839        assert!(matches!(
840            session.caught_up().await,
841            Err(CaughtUpError::Read(ReadSessionError::HeartbeatTimeout))
842        ));
843    }
844
845    #[tokio::test]
846    async fn dropping_session_rejects_wait() {
847        let caught_up = {
848            let session = test_session(stream::pending());
849            session.caught_up()
850        };
851
852        assert!(matches!(caught_up.await, Err(CaughtUpError::SessionClosed)));
853    }
854}
s2_sdk/session/read.rs

s2_sdk/session/
read.rs
