1use std::time::Duration;
2
3use futures::{Stream, StreamExt as _};
4use s2_common::{
5 basin::BasinName,
6 caps,
7 encryption::{EncryptionKey, EncryptionSpec},
8 read_extent::{EvaluatedReadLimit, ReadLimit, ReadUntil},
9 record::{Metered, MeteredSize as _, SeqNum, StreamPosition, Timestamp},
10 stream::{ReadEnd, ReadPosition, ReadSessionOutput, ReadStart, StreamName},
11};
12use s2_storage::record::{
13 StoredReadBatch, StoredReadSessionOutput, StoredSequencedRecord, decrypt_read_session_output,
14};
15use slatedb::config::{DurabilityLevel, ScanOptions};
16use tokio::{sync::broadcast, time::Instant};
17
18use super::{Backend, StreamHandle};
19use crate::{
20 backend::{
21 error::{
22 CheckTailError, ReadError, StorageError, StreamerMissingInActionError, UnwrittenError,
23 },
24 kv,
25 streamer::GuardedStreamerClient,
26 },
27 stream_id::StreamId,
28};
29
30impl Backend {
31 pub async fn open_for_check_tail(
32 &self,
33 basin: &BasinName,
34 stream: &StreamName,
35 ) -> Result<StreamHandle, CheckTailError> {
36 self.stream_handle_with_auto_create::<CheckTailError>(
37 basin,
38 stream,
39 |config| config.create_stream_on_read,
40 |_| Ok(EncryptionSpec::Plain),
41 )
42 .await
43 }
44
45 pub async fn open_for_read(
46 &self,
47 basin: &BasinName,
48 stream: &StreamName,
49 encryption_key: Option<EncryptionKey>,
50 ) -> Result<StreamHandle, ReadError> {
51 self.stream_handle_with_auto_create::<ReadError>(
52 basin,
53 stream,
54 |config| config.create_stream_on_read,
55 |cipher| Ok(EncryptionSpec::resolve(cipher, encryption_key)?),
56 )
57 .await
58 }
59}
60
61impl StreamHandle {
62 pub async fn check_tail(self) -> Result<StreamPosition, CheckTailError> {
63 let tail = self.client.check_tail().await?;
64 Ok(tail)
65 }
66
67 pub async fn read(
68 self,
69 start: ReadStart,
70 end: ReadEnd,
71 ) -> Result<impl Stream<Item = Result<ReadSessionOutput, ReadError>> + 'static, ReadError> {
72 let stream_id = self.client.stream_id();
73 let session = read_session(self.db, self.client, start, end).await?;
74 Ok(async_stream::stream! {
75 tokio::pin!(session);
76 while let Some(output) = session.next().await {
77 let output = match output {
78 Ok(output) => {
79 decrypt_read_session_output(output, &self.encryption, stream_id.as_bytes())
80 .map_err(ReadError::from)
81 }
82 Err(err) => Err(err),
83 };
84 let should_stop = output.is_err();
85 yield output;
86 if should_stop {
87 break;
88 }
89 }
90 })
91 }
92}
93
94async fn read_session(
95 db: slatedb::Db,
96 client: GuardedStreamerClient,
97 start: ReadStart,
98 end: ReadEnd,
99) -> Result<impl Stream<Item = Result<StoredReadSessionOutput, ReadError>> + 'static, ReadError> {
100 let stream_id = client.stream_id();
101 let tail = client.check_tail().await?;
102 let mut state = ReadSessionState {
103 start_seq_num: read_start_seq_num(&db, stream_id, start, end, tail).await?,
104 limit: EvaluatedReadLimit::Remaining(end.limit),
105 until: end.until,
106 wait: end.wait,
107 wait_deadline: None,
108 tail,
109 };
110 let session = async_stream::try_stream! {
111 'session: while let EvaluatedReadLimit::Remaining(limit) = state.limit {
112 if state.start_seq_num < state.tail.seq_num {
113 let prefix = kv::stream_record_data::ser_key_prefix(stream_id);
114 let start_suffix = kv::stream_record_data::ser_key_suffix(StreamPosition {
115 seq_num: state.start_seq_num,
116 timestamp: 0,
117 });
118 let end_suffix = kv::stream_record_data::ser_key_suffix(StreamPosition {
119 seq_num: state.tail.seq_num,
120 timestamp: 0,
121 });
122 let scan_opts = ScanOptions {
123 durability_filter: DurabilityLevel::Remote,
124 read_ahead_bytes: 1024 * 1024,
125 cache_blocks: true,
126 max_fetch_tasks: 8,
127 ..Default::default()
128 };
129 let mut it = db
130 .scan_prefix_with_options(prefix, start_suffix..end_suffix, &scan_opts)
131 .await?;
132
133 let mut records = Metered::with_capacity(
134 limit.count()
135 .unwrap_or(usize::MAX)
136 .min(caps::RECORD_BATCH_MAX.count),
137 );
138
139 while let EvaluatedReadLimit::Remaining(limit) = state.limit {
140 let Some(kv) = it.next().await? else {
141 break;
142 };
143 let (deser_stream_id, pos) = kv::stream_record_data::deser_key(kv.key)?;
144 assert_eq!(deser_stream_id, stream_id);
145
146 let record = kv::stream_record_data::deser_value(kv.value)?.sequenced(pos);
147
148 if end.until.deny(pos.timestamp)
149 || limit.deny(records.len() + 1, records.metered_size() + record.metered_size())
150 {
151 if records.is_empty() {
152 break 'session;
153 } else {
154 break;
155 }
156 }
157
158 if records.len() == caps::RECORD_BATCH_MAX.count
159 || records.metered_size() + record.metered_size() > caps::RECORD_BATCH_MAX.bytes
160 {
161 let new_records_buf = Metered::with_capacity(
162 limit.count()
163 .map_or(usize::MAX, |n| n.saturating_sub(records.len()))
164 .min(caps::RECORD_BATCH_MAX.count),
165 );
166 yield state.on_batch(StoredReadBatch {
167 records: std::mem::replace(&mut records, new_records_buf),
168 tail: None,
169 });
170 }
171
172 records.push(record);
173 }
174
175 if !records.is_empty() {
176 yield state.on_batch(StoredReadBatch {
177 records,
178 tail: None,
179 });
180 } else {
181 state.start_seq_num = state.tail.seq_num;
182 }
183 } else {
184 assert_eq!(state.start_seq_num, state.tail.seq_num);
185 if !end.may_follow() {
186 break;
187 }
188 match client.follow(state.start_seq_num).await? {
189 Ok(mut follow_rx) => {
190 state.arm_wait_deadline_if_unset();
192 if state.wait_deadline_expired() {
193 break;
194 }
195 yield StoredReadSessionOutput::Heartbeat(state.tail);
196 while let EvaluatedReadLimit::Remaining(limit) = state.limit {
197 tokio::select! {
198 biased;
199 msg = follow_rx.recv() => {
200 match msg {
201 Ok(mut records) => {
202 let count = records.len();
203 let tail = super::streamer::next_pos(&records);
204 let allowed_count = count_allowed_records(limit, end.until, &records);
205 if allowed_count > 0 {
206 yield state.on_batch(StoredReadBatch {
207 records: records.drain(..allowed_count).collect(),
208 tail: Some(tail),
209 });
210 }
211 if allowed_count < count {
212 break 'session;
213 }
214 Ok(())
215 }
216 Err(broadcast::error::RecvError::Lagged(_)) => {
217 continue 'session;
219 }
220 Err(broadcast::error::RecvError::Closed) => {
221 Err(StreamerMissingInActionError)
222 }
223 }
224 }
225 _ = new_heartbeat_sleep() => {
226 yield StoredReadSessionOutput::Heartbeat(state.tail);
227 Ok(())
228 }
229 _ = wait_sleep_until(state.wait_deadline) => {
230 break 'session;
231 }
232 }?;
233 }
234 }
235 Err(tail) => {
236 assert!(state.tail.seq_num < tail.seq_num, "tail cannot regress");
237 state.tail = tail;
238 }
239 }
240 }
241 }
242 };
243 Ok(session)
244}
245
246async fn read_start_seq_num(
247 db: &slatedb::Db,
248 stream_id: StreamId,
249 start: ReadStart,
250 end: ReadEnd,
251 tail: StreamPosition,
252) -> Result<SeqNum, ReadError> {
253 let mut read_pos = match start.from {
254 s2_common::stream::ReadFrom::SeqNum(seq_num) => ReadPosition::SeqNum(seq_num),
255 s2_common::stream::ReadFrom::Timestamp(timestamp) => ReadPosition::Timestamp(timestamp),
256 s2_common::stream::ReadFrom::TailOffset(tail_offset) => {
257 ReadPosition::SeqNum(tail.seq_num.saturating_sub(tail_offset))
258 }
259 };
260 if match read_pos {
261 ReadPosition::SeqNum(start_seq_num) => start_seq_num > tail.seq_num,
262 ReadPosition::Timestamp(start_timestamp) => start_timestamp > tail.timestamp,
263 } {
264 if start.clamp {
265 read_pos = ReadPosition::SeqNum(tail.seq_num);
266 } else {
267 return Err(UnwrittenError(tail).into());
268 }
269 }
270 if let ReadPosition::SeqNum(start_seq_num) = read_pos
271 && start_seq_num == tail.seq_num
272 && !end.may_follow()
273 {
274 return Err(UnwrittenError(tail).into());
275 }
276 Ok(match read_pos {
277 ReadPosition::SeqNum(start_seq_num) => start_seq_num,
278 ReadPosition::Timestamp(start_timestamp) => {
279 resolve_timestamp(db, stream_id, start_timestamp)
280 .await?
281 .unwrap_or(tail)
282 .seq_num
283 }
284 })
285}
286
287async fn resolve_timestamp(
288 db: &slatedb::Db,
289 stream_id: StreamId,
290 timestamp: Timestamp,
291) -> Result<Option<StreamPosition>, StorageError> {
292 let prefix = kv::stream_record_timestamp::ser_key_prefix(stream_id);
293 let start_suffix = kv::stream_record_timestamp::ser_key_suffix(StreamPosition {
294 seq_num: SeqNum::MIN,
295 timestamp,
296 });
297 let scan_opts = ScanOptions {
298 durability_filter: DurabilityLevel::Remote,
299 ..Default::default()
300 };
301 let mut it = db
302 .scan_prefix_with_options(prefix, start_suffix.., &scan_opts)
303 .await?;
304 Ok(match it.next().await? {
305 Some(kv) => {
306 let (deser_stream_id, pos) = kv::stream_record_timestamp::deser_key(kv.key)?;
307 assert_eq!(deser_stream_id, stream_id);
308 assert!(pos.timestamp >= timestamp);
309 kv::stream_record_timestamp::deser_value(kv.value)?;
310 Some(StreamPosition {
311 seq_num: pos.seq_num,
312 timestamp: pos.timestamp,
313 })
314 }
315 None => None,
316 })
317}
318
319struct ReadSessionState {
320 start_seq_num: u64,
321 limit: EvaluatedReadLimit,
322 until: ReadUntil,
323 wait: Option<Duration>,
324 wait_deadline: Option<Instant>,
325 tail: StreamPosition,
326}
327
328impl ReadSessionState {
329 fn arm_wait_deadline_if_unset(&mut self) {
330 if self.wait_deadline.is_none() {
331 self.reset_wait_deadline();
332 }
333 }
334
335 fn reset_wait_deadline(&mut self) {
336 self.wait_deadline = self.wait.map(|wait| Instant::now() + wait);
337 }
338
339 fn wait_deadline_expired(&self) -> bool {
340 self.wait_deadline
341 .is_some_and(|deadline| deadline <= Instant::now())
342 }
343
344 fn on_batch(&mut self, batch: StoredReadBatch) -> StoredReadSessionOutput {
345 if let Some(tail) = batch.tail {
346 self.tail = tail;
347 }
348 let last_record = batch.records.last().expect("non-empty");
349 let EvaluatedReadLimit::Remaining(limit) = self.limit else {
350 panic!("batch after exhausted limit");
351 };
352 let count = batch.records.len();
353 let bytes = batch.records.metered_size();
354 let last_position = last_record.position();
355 assert!(limit.allow(count, bytes));
356 assert!(self.until.allow(last_position.timestamp));
357 self.start_seq_num = last_position.seq_num + 1;
358 self.limit = limit.remaining(count, bytes);
359 self.reset_wait_deadline();
360 StoredReadSessionOutput::Batch(batch)
361 }
362}
363
364fn count_allowed_records(
365 limit: ReadLimit,
366 until: ReadUntil,
367 records: &[Metered<StoredSequencedRecord>],
368) -> usize {
369 let mut acc_size = 0;
370 let mut acc_count = 0;
371 for record in records {
372 if limit.deny(acc_count + 1, acc_size + record.metered_size())
373 || until.deny(record.position().timestamp)
374 {
375 break;
376 }
377 acc_count += 1;
378 acc_size += record.metered_size();
379 }
380 acc_count
381}
382
383#[cfg(not(test))]
384fn new_heartbeat_sleep() -> tokio::time::Sleep {
385 tokio::time::sleep(Duration::from_millis(rand::random_range(5_000..15_000)))
386}
387
388#[cfg(test)]
389fn new_heartbeat_sleep() -> tokio::time::Sleep {
390 tokio::time::sleep(Duration::from_millis(rand::random_range(5..15)))
391}
392
393async fn wait_sleep_until(deadline: Option<Instant>) {
394 match deadline {
395 Some(deadline) => tokio::time::sleep_until(deadline).await,
396 None => {
397 std::future::pending::<()>().await;
398 }
399 }
400}
401
402#[cfg(test)]
403mod tests {
404 use std::{sync::Arc, task::Poll};
405
406 use bytesize::ByteSize;
407 use futures::StreamExt;
408 use s2_common::{
409 basin::BasinName,
410 config::{BasinConfig, OptionalStreamConfig},
411 read_extent::{ReadLimit, ReadUntil},
412 record::{Metered, Record},
413 resources::ProvisionMode,
414 stream::{
415 AppendInput, AppendRecord, AppendRecordBatch, AppendRecordParts, ReadEnd, ReadFrom,
416 ReadSessionOutput, ReadStart, StreamName,
417 },
418 };
419 use slatedb::{Db, WriteBatch, object_store::memory::InMemory};
420 use tokio::time::Instant;
421
422 use super::*;
423 use crate::{
424 backend::{FOLLOWER_MAX_LAG, kv, streamer::DORMANT_TIMEOUT},
425 stream_id::StreamId,
426 };
427
428 fn append_input(record: Record) -> AppendInput {
429 let record: AppendRecord = AppendRecordParts {
430 timestamp: None,
431 record: Metered::from(record),
432 }
433 .try_into()
434 .unwrap();
435 let records: AppendRecordBatch = vec![record].try_into().unwrap();
436 AppendInput {
437 records,
438 match_seq_num: None,
439 fencing_token: None,
440 }
441 }
442
443 fn map_test_output(
444 output: Option<Result<ReadSessionOutput, ReadError>>,
445 ) -> Option<ReadSessionOutput> {
446 match output {
447 Some(Ok(output)) => Some(output),
448 Some(Err(e)) => panic!("Read error: {e:?}"),
449 None => None,
450 }
451 }
452
453 async fn poll_next_after_advance<S>(
454 session: &mut std::pin::Pin<Box<S>>,
455 advance_by: Duration,
456 ) -> Poll<Option<ReadSessionOutput>>
457 where
458 S: futures::Stream<Item = Result<ReadSessionOutput, ReadError>>,
459 {
460 let mut pinned_session = session.as_mut();
461 let next = pinned_session.next();
462 tokio::pin!(next);
463
464 assert!(
465 matches!(futures::poll!(&mut next), Poll::Pending),
466 "session unexpectedly yielded before time advanced"
467 );
468
469 tokio::time::advance(advance_by).await;
470 tokio::task::yield_now().await;
471
472 match futures::poll!(&mut next) {
473 Poll::Ready(output) => Poll::Ready(map_test_output(output)),
474 Poll::Pending => Poll::Pending,
475 }
476 }
477
478 #[tokio::test]
479 async fn resolve_timestamp_bounded_to_stream() {
480 let object_store = Arc::new(InMemory::new());
481 let db = Db::builder("/test", object_store).build().await.unwrap();
482 let backend = Backend::new(db, ByteSize::mib(10));
483
484 let stream_a: StreamId = [0u8; 32].into();
485 let stream_b: StreamId = [1u8; 32].into();
486
487 backend
488 .db
489 .put(
490 kv::stream_record_timestamp::ser_key(
491 stream_a,
492 StreamPosition {
493 seq_num: 0,
494 timestamp: 1000,
495 },
496 ),
497 kv::stream_record_timestamp::ser_value(),
498 )
499 .await
500 .unwrap();
501 backend
502 .db
503 .put(
504 kv::stream_record_timestamp::ser_key(
505 stream_b,
506 StreamPosition {
507 seq_num: 0,
508 timestamp: 2000,
509 },
510 ),
511 kv::stream_record_timestamp::ser_value(),
512 )
513 .await
514 .unwrap();
515
516 let result = resolve_timestamp(&backend.db, stream_a, 500).await.unwrap();
518 assert_eq!(
519 result,
520 Some(StreamPosition {
521 seq_num: 0,
522 timestamp: 1000
523 })
524 );
525
526 let result = resolve_timestamp(&backend.db, stream_a, 1500)
528 .await
529 .unwrap();
530 assert_eq!(result, None);
531 }
532
533 #[tokio::test]
534 async fn read_completes_when_all_records_deleted() {
535 let object_store = Arc::new(InMemory::new());
536 let db = Db::builder("/test", object_store).build().await.unwrap();
537 let backend = Backend::new(db, ByteSize::mib(10));
538
539 let basin: BasinName = "test-basin".parse().unwrap();
540 backend
541 .provision_basin(
542 basin.clone(),
543 BasinConfig::default(),
544 ProvisionMode::CreateOnly {
545 request_token: None,
546 },
547 )
548 .await
549 .unwrap();
550 let stream: StreamName = "test-stream".parse().unwrap();
551 backend
552 .provision_stream(
553 basin.clone(),
554 stream.clone(),
555 OptionalStreamConfig::default(),
556 ProvisionMode::CreateOnly {
557 request_token: None,
558 },
559 )
560 .await
561 .unwrap();
562
563 let input = append_input(Record::try_from_parts(vec![], bytes::Bytes::from("x")).unwrap());
564 let ack = backend
565 .open_for_append(&basin, &stream, None)
566 .await
567 .unwrap()
568 .append(input)
569 .await
570 .unwrap();
571 assert!(ack.end.seq_num > 0);
572
573 let stream_id = StreamId::new(&basin, &stream);
574 let mut batch = WriteBatch::new();
575 batch.delete(kv::stream_record_data::ser_key(stream_id, ack.start));
576 backend.db.write(batch).await.unwrap();
577
578 let start = ReadStart {
579 from: ReadFrom::SeqNum(0),
580 clamp: false,
581 };
582 let end = ReadEnd {
583 limit: ReadLimit::Count(10),
584 until: ReadUntil::Unbounded,
585 wait: None,
586 };
587 let session = backend
588 .open_for_read(&basin, &stream, None)
589 .await
590 .unwrap()
591 .read(start, end)
592 .await
593 .unwrap();
594 let records: Vec<_> = tokio::time::timeout(
595 Duration::from_secs(2),
596 futures::StreamExt::collect::<Vec<_>>(session),
597 )
598 .await
599 .expect("read should not spin forever");
600 assert!(records.into_iter().all(|r| r.is_ok()));
601 }
602
603 #[tokio::test(flavor = "current_thread", start_paused = true)]
604 async fn read_wait_is_not_extended_by_heartbeats() {
605 let object_store = Arc::new(InMemory::new());
606 let db = Db::builder("/test", object_store).build().await.unwrap();
607 let backend = Backend::new(db, ByteSize::mib(10));
608
609 let basin: BasinName = "test-basin".parse().unwrap();
610 backend
611 .provision_basin(
612 basin.clone(),
613 BasinConfig::default(),
614 ProvisionMode::CreateOnly {
615 request_token: None,
616 },
617 )
618 .await
619 .unwrap();
620 let stream: StreamName = "test-stream".parse().unwrap();
621 backend
622 .provision_stream(
623 basin.clone(),
624 stream.clone(),
625 OptionalStreamConfig::default(),
626 ProvisionMode::CreateOnly {
627 request_token: None,
628 },
629 )
630 .await
631 .unwrap();
632
633 let wait = Duration::from_millis(30);
634 let start = ReadStart {
635 from: ReadFrom::SeqNum(0),
636 clamp: false,
637 };
638 let end = ReadEnd {
639 limit: ReadLimit::Unbounded,
640 until: ReadUntil::Unbounded,
641 wait: Some(wait),
642 };
643
644 let session = backend
645 .open_for_read(&basin, &stream, None)
646 .await
647 .unwrap()
648 .read(start, end)
649 .await
650 .unwrap();
651 let mut session = Box::pin(session);
652 let probe_step = Duration::from_millis(1);
653 let first = session
654 .as_mut()
655 .next()
656 .await
657 .expect("session should enter follow mode")
658 .expect("session should not error");
659 assert!(matches!(first, ReadSessionOutput::Heartbeat(_)));
660
661 let started = Instant::now();
662 let second = match poll_next_after_advance(&mut session, wait).await {
663 Poll::Ready(Some(output)) => output,
664 Poll::Ready(None) => panic!("session closed before emitting a follow heartbeat"),
665 Poll::Pending => panic!("expected a follow heartbeat before the wait budget expired"),
666 };
667 assert!(matches!(second, ReadSessionOutput::Heartbeat(_)));
668
669 tokio::task::yield_now().await;
670 let closed_at = loop {
671 match futures::poll!(session.as_mut().next()) {
672 Poll::Ready(Some(Ok(ReadSessionOutput::Heartbeat(_)))) => {}
673 Poll::Ready(Some(Ok(output))) => {
674 panic!("unexpected output after wait deadline: {output:?}");
675 }
676 Poll::Ready(Some(Err(e))) => panic!("Read error: {e:?}"),
677 Poll::Ready(None) => break Instant::now(),
678 Poll::Pending => panic!("session should close once the wait budget expires"),
679 }
680 };
681
682 assert!(closed_at >= started + wait);
683 assert!(closed_at <= started + wait + probe_step);
684 }
685
686 #[tokio::test(flavor = "current_thread", start_paused = true)]
687 async fn read_wait_is_reset_by_delivered_follow_batch() {
688 let object_store = Arc::new(InMemory::new());
689 let db = Db::builder("/test", object_store).build().await.unwrap();
690 let backend = Backend::new(db, ByteSize::mib(10));
691
692 let basin: BasinName = "test-basin".parse().unwrap();
693 backend
694 .provision_basin(
695 basin.clone(),
696 BasinConfig::default(),
697 ProvisionMode::CreateOnly {
698 request_token: None,
699 },
700 )
701 .await
702 .unwrap();
703 let stream: StreamName = "test-stream".parse().unwrap();
704 backend
705 .provision_stream(
706 basin.clone(),
707 stream.clone(),
708 OptionalStreamConfig::default(),
709 ProvisionMode::CreateOnly {
710 request_token: None,
711 },
712 )
713 .await
714 .unwrap();
715
716 let initial_input =
717 append_input(Record::try_from_parts(vec![], bytes::Bytes::from("initial")).unwrap());
718 backend
719 .open_for_append(&basin, &stream, None)
720 .await
721 .unwrap()
722 .append(initial_input)
723 .await
724 .unwrap();
725
726 let wait = Duration::from_millis(30);
727 let probe_step = Duration::from_millis(1);
728 let start = ReadStart {
729 from: ReadFrom::SeqNum(0),
730 clamp: false,
731 };
732 let end = ReadEnd {
733 limit: ReadLimit::Unbounded,
734 until: ReadUntil::Unbounded,
735 wait: Some(wait),
736 };
737
738 let session = backend
739 .open_for_read(&basin, &stream, None)
740 .await
741 .unwrap()
742 .read(start, end)
743 .await
744 .unwrap();
745 let mut session = Box::pin(session);
746
747 let first = session
748 .as_mut()
749 .next()
750 .await
751 .expect("session should yield the initial batch")
752 .expect("session should not error");
753 let ReadSessionOutput::Batch(batch) = first else {
754 panic!("expected initial batch");
755 };
756 let initial_record = batch
757 .records
758 .first()
759 .expect("batch should contain one record");
760 let Record::Envelope(initial_envelope) = initial_record.inner() else {
761 panic!("expected plaintext envelope record");
762 };
763 assert_eq!(initial_envelope.body().as_ref(), b"initial");
764
765 let second = session
766 .as_mut()
767 .next()
768 .await
769 .expect("session should enter follow mode")
770 .expect("session should not error");
771 assert!(matches!(second, ReadSessionOutput::Heartbeat(_)));
772
773 tokio::time::advance(Duration::from_millis(20)).await;
774 tokio::task::yield_now().await;
775
776 let follow_input =
777 append_input(Record::try_from_parts(vec![], bytes::Bytes::from("follow-1")).unwrap());
778 backend
779 .open_for_append(&basin, &stream, None)
780 .await
781 .unwrap()
782 .append(follow_input)
783 .await
784 .unwrap();
785
786 let follow = session
787 .as_mut()
788 .next()
789 .await
790 .expect("session should deliver the live batch")
791 .expect("session should not error");
792 let reset_at = Instant::now();
793 let ReadSessionOutput::Batch(batch) = follow else {
794 panic!("expected live batch after append");
795 };
796 let follow_record = batch
797 .records
798 .first()
799 .expect("batch should contain one record");
800 let Record::Envelope(follow_envelope) = follow_record.inner() else {
801 panic!("expected plaintext envelope record");
802 };
803 assert_eq!(follow_envelope.body().as_ref(), b"follow-1");
804
805 tokio::time::advance(wait - probe_step).await;
806 tokio::task::yield_now().await;
807
808 loop {
809 match futures::poll!(session.as_mut().next()) {
810 Poll::Ready(Some(Ok(ReadSessionOutput::Heartbeat(_)))) => {}
811 Poll::Ready(Some(Ok(output))) => {
812 panic!("unexpected output before the reset wait deadline: {output:?}");
813 }
814 Poll::Ready(Some(Err(e))) => panic!("Read error: {e:?}"),
815 Poll::Ready(None) => {
816 panic!("session closed before the reset wait budget expired");
817 }
818 Poll::Pending => break,
819 }
820 }
821
822 tokio::time::advance(probe_step).await;
823 tokio::task::yield_now().await;
824
825 let closed_at = loop {
826 match futures::poll!(session.as_mut().next()) {
827 Poll::Ready(Some(Ok(ReadSessionOutput::Heartbeat(_)))) => {}
828 Poll::Ready(Some(Ok(output))) => {
829 panic!("unexpected output after the reset wait deadline: {output:?}");
830 }
831 Poll::Ready(Some(Err(e))) => panic!("Read error: {e:?}"),
832 Poll::Ready(None) => break Instant::now(),
833 Poll::Pending => {
834 panic!("session should close once the reset wait budget expires");
835 }
836 }
837 };
838
839 assert!(closed_at >= reset_at + wait);
840 assert!(closed_at <= reset_at + wait + probe_step);
841 }
842
843 #[tokio::test(flavor = "current_thread", start_paused = true)]
844 async fn read_wait_is_not_reset_after_follow_lag_without_catchup_records() {
845 let object_store = Arc::new(InMemory::new());
846 let db = Db::builder("/test", object_store).build().await.unwrap();
847 let backend = Backend::new(db, ByteSize::mib(10));
848
849 let basin: BasinName = "test-basin".parse().unwrap();
850 backend
851 .provision_basin(
852 basin.clone(),
853 BasinConfig::default(),
854 ProvisionMode::CreateOnly {
855 request_token: None,
856 },
857 )
858 .await
859 .unwrap();
860 let stream: StreamName = "test-stream".parse().unwrap();
861 backend
862 .provision_stream(
863 basin.clone(),
864 stream.clone(),
865 OptionalStreamConfig::default(),
866 ProvisionMode::CreateOnly {
867 request_token: None,
868 },
869 )
870 .await
871 .unwrap();
872
873 let wait = Duration::from_secs(30);
874 let start = ReadStart {
875 from: ReadFrom::SeqNum(0),
876 clamp: false,
877 };
878 let end = ReadEnd {
879 limit: ReadLimit::Unbounded,
880 until: ReadUntil::Unbounded,
881 wait: Some(wait),
882 };
883 let session = backend
884 .open_for_read(&basin, &stream, None)
885 .await
886 .unwrap()
887 .read(start, end)
888 .await
889 .unwrap();
890 let mut session = Box::pin(session);
891
892 let first = session
893 .as_mut()
894 .next()
895 .await
896 .expect("session should enter follow mode")
897 .expect("session should not error");
898 assert!(matches!(first, ReadSessionOutput::Heartbeat(_)));
899
900 let stream_id = StreamId::new(&basin, &stream);
901 let mut delete_batch = WriteBatch::new();
902 let lagged_appends = FOLLOWER_MAX_LAG + 25;
903
904 for i in 0..lagged_appends {
905 let input = append_input(
906 Record::try_from_parts(vec![], bytes::Bytes::from(format!("lagged-{i}"))).unwrap(),
907 );
908 let ack = backend
909 .open_for_append(&basin, &stream, None)
910 .await
911 .unwrap()
912 .append(input)
913 .await
914 .unwrap();
915 delete_batch.delete(kv::stream_record_data::ser_key(stream_id, ack.start));
916 }
917
918 backend.db.write(delete_batch).await.unwrap();
919
920 tokio::time::advance(wait + Duration::from_secs(1)).await;
921 tokio::task::yield_now().await;
922
923 let next = session.as_mut().next().await;
924 assert!(
925 next.is_none(),
926 "session should close immediately once the original wait budget has elapsed"
927 );
928 }
929
930 #[tokio::test(flavor = "current_thread", start_paused = true)]
931 async fn unbounded_follow_survives_streamer_dormancy() {
932 let object_store = Arc::new(InMemory::new());
933 let db = Db::builder("/test", object_store).build().await.unwrap();
934 let backend = Backend::new(db, ByteSize::mib(10));
935
936 let basin: BasinName = "test-basin".parse().unwrap();
937 backend
938 .provision_basin(
939 basin.clone(),
940 BasinConfig::default(),
941 ProvisionMode::CreateOnly {
942 request_token: None,
943 },
944 )
945 .await
946 .unwrap();
947 let stream: StreamName = "test-stream".parse().unwrap();
948 backend
949 .provision_stream(
950 basin.clone(),
951 stream.clone(),
952 OptionalStreamConfig::default(),
953 ProvisionMode::CreateOnly {
954 request_token: None,
955 },
956 )
957 .await
958 .unwrap();
959
960 let initial_input =
961 append_input(Record::try_from_parts(vec![], bytes::Bytes::from("initial")).unwrap());
962 backend
963 .open_for_append(&basin, &stream, None)
964 .await
965 .unwrap()
966 .append(initial_input)
967 .await
968 .unwrap();
969
970 let start = ReadStart {
971 from: ReadFrom::SeqNum(0),
972 clamp: false,
973 };
974 let end = ReadEnd {
975 limit: ReadLimit::Unbounded,
976 until: ReadUntil::Unbounded,
977 wait: None,
978 };
979 let session = backend
980 .open_for_read(&basin, &stream, None)
981 .await
982 .unwrap()
983 .read(start, end)
984 .await
985 .unwrap();
986 let mut session = Box::pin(session);
987
988 let first = session
989 .as_mut()
990 .next()
991 .await
992 .expect("session should yield initial batch")
993 .expect("session should not error");
994 assert!(matches!(first, ReadSessionOutput::Batch(_)));
995
996 let second = session
997 .as_mut()
998 .next()
999 .await
1000 .expect("session should enter follow mode")
1001 .expect("session should not error");
1002 assert!(matches!(second, ReadSessionOutput::Heartbeat(_)));
1003
1004 tokio::time::advance(DORMANT_TIMEOUT + Duration::from_secs(1)).await;
1005 tokio::task::yield_now().await;
1006
1007 let follow_input =
1008 append_input(Record::try_from_parts(vec![], bytes::Bytes::from("follow-1")).unwrap());
1009 backend
1010 .open_for_append(&basin, &stream, None)
1011 .await
1012 .unwrap()
1013 .append(follow_input)
1014 .await
1015 .unwrap();
1016
1017 let next = session
1018 .as_mut()
1019 .next()
1020 .await
1021 .expect("session should stay open after dormancy")
1022 .expect("session should not error after dormancy");
1023 let ReadSessionOutput::Batch(batch) = next else {
1024 panic!("expected new batch after append");
1025 };
1026 assert_eq!(batch.records.len(), 1);
1027 let record = batch.records.first().expect("batch should have one record");
1028 let Record::Envelope(envelope) = record.inner() else {
1029 panic!("expected envelope record");
1030 };
1031 assert_eq!(envelope.body().as_ref(), b"follow-1");
1032 }
1033}