1use std::collections::HashMap;
10use std::sync::Arc;
11use std::time::Duration;
12
13use std::sync::atomic::{AtomicU64, Ordering};
14use std::sync::{Mutex, MutexGuard};
15
16use aws_sdk_kinesis::operation::get_records::GetRecordsError;
17use aws_sdk_kinesis::operation::get_shard_iterator::builders::GetShardIteratorFluentBuilder;
18use aws_sdk_kinesis::primitives::DateTime;
19use aws_sdk_kinesis::types::{Shard, ShardIteratorType};
20use futures::Stream;
21use ruststream::Subscriber;
22use tokio::sync::{mpsc, oneshot};
23
24use crate::broker::Core;
25use crate::error::{KinesisError, sdk_err};
26use crate::lease::{LeaseStore, SHARD_END};
27use crate::message::{KPL_MAGIC, KinesisMessage, KinesisPosition, Settlement};
28use crate::stream::KinesisStream;
29use crate::track::Watermark;
30
31const SHARD_SYNC: Duration = Duration::from_secs(10);
33const LEASE_TTL: Duration = Duration::from_secs(10);
35const RENEW_EVERY: Duration = Duration::from_secs(3);
36const CHANNEL_CAPACITY: usize = 64;
38
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub(crate) enum StreamStart {
44 Horizon,
45 Latest,
46 Timestamp(u64),
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
51pub(crate) enum ShardStart {
52 Stream(StreamStart),
54 At(String),
56 After(String),
58}
59
60pub(crate) struct ShardSeek {
62 pub(crate) start: ShardStart,
63 pub(crate) done: oneshot::Sender<Result<(), KinesisError>>,
64}
65
66#[derive(Clone)]
69pub(crate) struct ShardHandle {
70 pub(crate) gate: Arc<AtomicU64>,
71 pub(crate) tx: mpsc::UnboundedSender<ShardSeek>,
72}
73
74#[derive(Default)]
76pub(crate) struct SeekState {
77 shards: Mutex<HashMap<String, ShardHandle>>,
79 start: Mutex<Option<StreamStart>>,
84}
85
86pub(crate) type SeekBus = Arc<SeekState>;
87
88impl SeekState {
89 fn register(&self, shard: String, handle: ShardHandle) {
90 self.lock_shards().insert(shard, handle);
91 }
92
93 fn deregister(&self, shard: &str) {
94 self.lock_shards().remove(shard);
95 }
96
97 fn handle(&self, shard: &str) -> Option<ShardHandle> {
98 self.lock_shards().get(shard).cloned()
99 }
100
101 fn live(&self) -> Vec<ShardHandle> {
102 self.lock_shards().values().cloned().collect()
103 }
104
105 fn install(&self, start: StreamStart) {
106 *self.start.lock().expect("seek state mutex poisoned") = Some(start);
107 }
108
109 pub(crate) fn installed(&self) -> Option<StreamStart> {
110 *self.start.lock().expect("seek state mutex poisoned")
111 }
112
113 fn lock_shards(&self) -> MutexGuard<'_, HashMap<String, ShardHandle>> {
114 self.shards.lock().expect("seek state mutex poisoned")
115 }
116}
117
118pub(crate) struct Stamped {
121 stamp: Option<(u64, Arc<AtomicU64>)>,
122 item: Result<KinesisMessage, KinesisError>,
123}
124
125impl Stamped {
126 fn live(epoch: u64, gate: &Arc<AtomicU64>, item: Result<KinesisMessage, KinesisError>) -> Self {
127 Self {
128 stamp: Some((epoch, Arc::clone(gate))),
129 item,
130 }
131 }
132
133 fn unstamped(item: Result<KinesisMessage, KinesisError>) -> Self {
134 Self { stamp: None, item }
135 }
136
137 fn current(&self) -> bool {
138 self.stamp
139 .as_ref()
140 .is_none_or(|(epoch, gate)| *epoch == gate.load(Ordering::Acquire))
141 }
142}
143
144pub struct KinesisSubscriber {
149 stream: String,
150 rx: mpsc::Receiver<Stamped>,
151 bus: SeekBus,
152}
153
154impl std::fmt::Debug for KinesisSubscriber {
155 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
156 f.debug_struct("KinesisSubscriber")
157 .field("stream", &self.stream)
158 .finish_non_exhaustive()
159 }
160}
161
162impl KinesisSubscriber {
163 #[must_use]
165 pub fn stream_name(&self) -> &str {
166 &self.stream
167 }
168
169 pub(crate) fn open(core: &Core, descriptor: KinesisStream) -> Self {
170 let (tx, rx) = mpsc::channel(CHANNEL_CAPACITY);
171 let stream = descriptor.stream().to_owned();
172 let bus: SeekBus = Arc::new(SeekState::default());
173 tokio::spawn(coordinate(
174 core.client.clone(),
175 Arc::clone(&core.store),
176 core.owner.clone(),
177 descriptor,
178 tx,
179 Arc::clone(&bus),
180 ));
181 Self { stream, rx, bus }
182 }
183}
184
185#[derive(Clone)]
195pub struct KinesisSeeker {
196 bus: SeekBus,
197}
198
199impl std::fmt::Debug for KinesisSeeker {
200 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
201 f.debug_struct("KinesisSeeker").finish_non_exhaustive()
202 }
203}
204
205impl KinesisSeeker {
206 async fn seek_shard(&self, shard: String, start: ShardStart) -> Result<(), KinesisError> {
208 let Some(handle) = self.bus.handle(&shard) else {
209 return Err(KinesisError::Read {
210 stream: String::new(),
211 shard,
212 source: Box::from("no live reader for this shard (not owned, or finished)"),
213 });
214 };
215 handle.gate.fetch_add(1, Ordering::Release);
218 let (done, wait) = oneshot::channel();
219 handle
220 .tx
221 .send(ShardSeek { start, done })
222 .map_err(|_| KinesisError::Read {
223 stream: String::new(),
224 shard: shard.clone(),
225 source: Box::from("the shard reader has shut down"),
226 })?;
227 wait.await.map_err(|_| KinesisError::Read {
228 stream: String::new(),
229 shard,
230 source: Box::from("the shard reader has shut down"),
231 })?
232 }
233
234 async fn seek_stream(&self, start: StreamStart) -> Result<(), KinesisError> {
236 self.bus.install(start);
241 let handles = self.bus.live();
242 for handle in &handles {
245 handle.gate.fetch_add(1, Ordering::Release);
246 }
247 let mut pending = Vec::with_capacity(handles.len());
248 for handle in handles {
249 let (done, wait) = oneshot::channel();
250 if handle
253 .tx
254 .send(ShardSeek {
255 start: ShardStart::Stream(start),
256 done,
257 })
258 .is_ok()
259 {
260 pending.push(wait);
261 }
262 }
263 for wait in pending {
264 if let Ok(outcome) = wait.await {
265 outcome?;
266 }
267 }
268 Ok(())
269 }
270}
271
272impl ruststream::Seeker for KinesisSeeker {
273 type Position = KinesisPosition;
274 type Error = KinesisError;
275
276 async fn seek(&self, to: KinesisPosition) -> Result<(), KinesisError> {
277 match to {
278 KinesisPosition::Horizon => self.seek_stream(StreamStart::Horizon).await,
279 KinesisPosition::Latest => self.seek_stream(StreamStart::Latest).await,
280 KinesisPosition::Timestamp(millis) => {
281 self.seek_stream(StreamStart::Timestamp(millis)).await
282 }
283 KinesisPosition::Sequence { shard, sequence } => {
284 self.seek_shard(shard, ShardStart::At(sequence)).await
285 }
286 }
287 }
288}
289
290impl ruststream::Seekable for KinesisSubscriber {
291 type Seeker = KinesisSeeker;
292
293 fn seeker(&self) -> KinesisSeeker {
294 KinesisSeeker {
295 bus: Arc::clone(&self.bus),
296 }
297 }
298}
299
300impl Subscriber for KinesisSubscriber {
301 type Message = KinesisMessage;
302 type Error = KinesisError;
303
304 fn stream(&mut self) -> impl Stream<Item = Result<KinesisMessage, KinesisError>> + Send + '_ {
305 futures::stream::poll_fn(move |cx| {
310 loop {
311 match self.rx.poll_recv(cx) {
312 std::task::Poll::Ready(Some(stamped)) => {
313 if stamped.current() {
314 return std::task::Poll::Ready(Some(stamped.item));
315 }
316 }
317 std::task::Poll::Ready(None) => return std::task::Poll::Ready(None),
318 std::task::Poll::Pending => return std::task::Poll::Pending,
319 }
320 }
321 })
322 }
323}
324
325async fn list_all_shards(
326 client: &aws_sdk_kinesis::Client,
327 stream: &str,
328) -> Result<Vec<Shard>, KinesisError> {
329 let mut shards = Vec::new();
331 let mut token: Option<String> = None;
332 loop {
333 let request = token.take().map_or_else(
334 || client.list_shards().stream_name(stream),
335 |t| client.list_shards().next_token(t),
336 );
337 let output = request.send().await.map_err(|e| KinesisError::Stream {
338 stream: stream.to_owned(),
339 source: sdk_err(&e),
340 })?;
341 shards.extend(output.shards().iter().cloned());
342 match output.next_token() {
343 Some(t) => token = Some(t.to_owned()),
344 None => return Ok(shards),
345 }
346 }
347}
348
349fn is_closed(shard: &Shard) -> bool {
350 shard
351 .sequence_number_range()
352 .and_then(|r| r.ending_sequence_number())
353 .is_some()
354}
355
356async fn coordinate(
357 client: aws_sdk_kinesis::Client,
358 store: Arc<dyn LeaseStore>,
359 owner: String,
360 descriptor: KinesisStream,
361 out: mpsc::Sender<Stamped>,
362 bus: SeekBus,
363) {
364 let stream = descriptor.stream().to_owned();
365 let mut readers: HashMap<String, tokio::task::JoinHandle<()>> = HashMap::new();
366 loop {
367 readers.retain(|_, handle| !handle.is_finished());
368
369 match list_all_shards(&client, &stream).await {
370 Ok(shards) => {
371 let by_id: HashMap<&str, &Shard> =
372 shards.iter().map(|s| (s.shard_id(), s)).collect();
373 for shard in &shards {
374 let id = shard.shard_id().to_owned();
375 if readers.contains_key(&id) {
376 continue;
377 }
378 if !parents_done(&bus, shard, &by_id, store.as_ref()).await {
379 continue;
380 }
381 match store.read(&id).await {
382 Ok(state) if state.checkpoint.as_deref() == Some(SHARD_END) => continue,
383 Ok(_) => {}
384 Err(err) => {
385 let _ = out
386 .send(Stamped::unstamped(Err(KinesisError::Lease {
387 shard: id.clone(),
388 source: err,
389 })))
390 .await;
391 continue;
392 }
393 }
394 match store.acquire(&id, &owner, LEASE_TTL).await {
395 Ok(true) => {
396 let (seek_tx, seek_rx) = mpsc::unbounded_channel();
397 let handle = ShardHandle {
398 gate: Arc::new(AtomicU64::new(0)),
399 tx: seek_tx,
400 };
401 bus.register(id.clone(), handle.clone());
402 readers.insert(
403 id.clone(),
404 tokio::spawn(read_shard(
405 client.clone(),
406 Arc::clone(&store),
407 owner.clone(),
408 descriptor.clone(),
409 id,
410 out.clone(),
411 handle.gate,
412 seek_rx,
413 Arc::clone(&bus),
414 )),
415 );
416 }
417 Ok(false) => {} Err(err) => {
419 let _ = out
420 .send(Stamped::unstamped(Err(KinesisError::Lease {
421 shard: id.clone(),
422 source: err,
423 })))
424 .await;
425 }
426 }
427 }
428 }
429 Err(err) => {
430 if out.send(Stamped::unstamped(Err(err))).await.is_err() {
431 break;
432 }
433 }
434 }
435
436 tokio::select! {
437 () = out.closed() => break,
438 () = tokio::time::sleep(SHARD_SYNC) => {}
439 }
440 }
441 }
443
444async fn parents_done(
449 bus: &SeekState,
450 shard: &Shard,
451 by_id: &HashMap<&str, &Shard>,
452 store: &dyn LeaseStore,
453) -> bool {
454 let parents = [shard.parent_shard_id(), shard.adjacent_parent_shard_id()];
455 for parent in parents.into_iter().flatten() {
456 let Some(parent_shard) = by_id.get(parent) else {
457 continue; };
459 let Ok(state) = store.read(parent).await else {
460 return false;
461 };
462 match state.checkpoint.as_deref() {
463 Some(SHARD_END) => {}
464 None if is_closed(parent_shard)
465 && matches!(bus.installed(), None | Some(StreamStart::Latest)) => {}
466 _ => return false,
467 }
468 }
469 true
470}
471
472fn iterator_at(
474 request: GetShardIteratorFluentBuilder,
475 start: &ShardStart,
476) -> GetShardIteratorFluentBuilder {
477 match start {
478 ShardStart::Stream(StreamStart::Latest) => {
479 request.shard_iterator_type(ShardIteratorType::Latest)
480 }
481 ShardStart::Stream(StreamStart::Horizon) => {
482 request.shard_iterator_type(ShardIteratorType::TrimHorizon)
483 }
484 ShardStart::Stream(StreamStart::Timestamp(millis)) => request
485 .shard_iterator_type(ShardIteratorType::AtTimestamp)
486 .timestamp(DateTime::from_millis(
487 i64::try_from(*millis).unwrap_or(i64::MAX),
488 )),
489 ShardStart::At(sequence) => request
490 .shard_iterator_type(ShardIteratorType::AtSequenceNumber)
491 .starting_sequence_number(sequence),
492 ShardStart::After(sequence) => request
493 .shard_iterator_type(ShardIteratorType::AfterSequenceNumber)
494 .starting_sequence_number(sequence),
495 }
496}
497
498#[derive(Debug, Clone, Copy, PartialEq, Eq)]
501enum Reopen {
502 Start,
505 Recover,
510}
511
512async fn initial_iterator(
513 client: &aws_sdk_kinesis::Client,
514 stream: &str,
515 shard: &str,
516 bus: &SeekState,
517 store: &dyn LeaseStore,
518 why: Reopen,
519) -> Result<Option<String>, KinesisError> {
520 let checkpoint = store
521 .read(shard)
522 .await
523 .map_err(|e| KinesisError::Lease {
524 shard: shard.to_owned(),
525 source: e,
526 })?
527 .checkpoint;
528 if checkpoint.as_deref() == Some(SHARD_END) {
529 return Ok(None);
530 }
531 let start = match (why, bus.installed(), checkpoint) {
532 (Reopen::Start, Some(installed), _) | (Reopen::Recover, Some(installed), None) => {
533 ShardStart::Stream(installed)
534 }
535 (_, _, Some(sequence)) => ShardStart::After(sequence),
536 (_, None, None) => ShardStart::Stream(StreamStart::Latest),
537 };
538 let request = iterator_at(
539 client
540 .get_shard_iterator()
541 .stream_name(stream)
542 .shard_id(shard),
543 &start,
544 );
545 let output = request.send().await.map_err(|e| KinesisError::Read {
546 stream: stream.to_owned(),
547 shard: shard.to_owned(),
548 source: sdk_err(&e),
549 })?;
550 Ok(output.shard_iterator().map(str::to_owned))
551}
552
553async fn apply_shard_seek(
555 client: &aws_sdk_kinesis::Client,
556 stream: &str,
557 shard: &str,
558 seek: ShardSeek,
559 iterator: &mut String,
560 tracker: &mut Arc<Watermark>,
561) {
562 let fresh = iterator_at(
563 client
564 .get_shard_iterator()
565 .stream_name(stream)
566 .shard_id(shard),
567 &seek.start,
568 )
569 .send()
570 .await;
571 let outcome = match fresh {
572 Ok(output) => output.shard_iterator().map_or_else(
573 || {
574 Err(KinesisError::Read {
575 stream: stream.to_owned(),
576 shard: shard.to_owned(),
577 source: Box::from("the service returned no iterator for the position"),
578 })
579 },
580 |new_iterator| {
581 new_iterator.clone_into(iterator);
582 *tracker = Arc::new(Watermark::default());
583 Ok(())
584 },
585 ),
586 Err(err) => Err(KinesisError::Read {
587 stream: stream.to_owned(),
588 shard: shard.to_owned(),
589 source: sdk_err(&err),
590 }),
591 };
592 let _ = seek.done.send(outcome);
593}
594
595#[allow(clippy::too_many_lines, clippy::too_many_arguments)]
596async fn read_shard(
597 client: aws_sdk_kinesis::Client,
598 store: Arc<dyn LeaseStore>,
599 owner: String,
600 descriptor: KinesisStream,
601 shard: String,
602 out: mpsc::Sender<Stamped>,
603 gate: Arc<AtomicU64>,
604 mut seek_rx: mpsc::UnboundedReceiver<ShardSeek>,
605 bus: SeekBus,
606) {
607 struct BusGuard {
609 bus: SeekBus,
610 shard: String,
611 }
612 impl Drop for BusGuard {
613 fn drop(&mut self) {
614 self.bus.deregister(&self.shard);
615 }
616 }
617 let _bus_guard = BusGuard {
618 bus: Arc::clone(&bus),
619 shard: shard.clone(),
620 };
621 let stream = descriptor.stream().to_owned();
622 let mut tracker = Arc::new(Watermark::default());
623 let mut iterator = match initial_iterator(
624 &client,
625 &stream,
626 &shard,
627 &bus,
628 store.as_ref(),
629 Reopen::Start,
630 )
631 .await
632 {
633 Ok(Some(iterator)) => iterator,
634 Ok(None) => return, Err(err) => {
636 let _ = out.send(Stamped::unstamped(Err(err))).await;
637 return;
638 }
639 };
640 let mut last_renew = tokio::time::Instant::now();
641 let mut failures: u32 = 0;
642
643 loop {
644 if out.is_closed() {
645 let _ = store.release(&shard, &owner).await;
646 return;
647 }
648 while let Ok(seek) = seek_rx.try_recv() {
652 apply_shard_seek(&client, &stream, &shard, seek, &mut iterator, &mut tracker).await;
653 }
654 if last_renew.elapsed() >= RENEW_EVERY {
655 match store.renew(&shard, &owner, LEASE_TTL).await {
656 Ok(true) => last_renew = tokio::time::Instant::now(),
657 Ok(false) | Err(_) => return,
660 }
661 }
662
663 let epoch = gate.load(Ordering::Acquire);
666 let response = client
667 .get_records()
668 .shard_iterator(&iterator)
669 .limit(descriptor.batch_value())
670 .send()
671 .await;
672 match response {
673 Ok(output) => {
674 failures = 0;
675 for record in output.records() {
676 let data = record.data().as_ref();
677 if data.len() > 4 && data[0..4] == KPL_MAGIC {
678 if out
679 .send(Stamped::live(
680 epoch,
681 &gate,
682 Err(KinesisError::AggregatedRecord {
683 shard: shard.clone(),
684 }),
685 ))
686 .await
687 .is_err()
688 {
689 return;
690 }
691 continue;
692 }
693 let settlement = Settlement {
694 tracker: Arc::clone(&tracker),
695 index: tracker.deliver(record.sequence_number()),
696 store: Arc::clone(&store),
697 shard: shard.clone(),
698 owner: owner.clone(),
699 epoch,
700 gate: Arc::clone(&gate),
701 };
702 let message = KinesisMessage::new(
703 data,
704 record.partition_key(),
705 record.sequence_number(),
706 settlement,
707 );
708 if out
709 .send(Stamped::live(epoch, &gate, Ok(message)))
710 .await
711 .is_err()
712 {
713 let _ = store.release(&shard, &owner).await;
714 return;
715 }
716 }
717
718 let Some(next) = output.next_shard_iterator() else {
719 while !tracker.drained() && !out.is_closed() {
722 tokio::time::sleep(Duration::from_millis(100)).await;
723 }
724 if tracker.drained() {
725 let _ = store.checkpoint(&shard, &owner, SHARD_END).await;
726 }
727 let _ = store.release(&shard, &owner).await;
728 return;
729 };
730 iterator = next.to_owned();
731 if output.records().is_empty() {
732 tokio::select! {
733 () = out.closed() => {}
734 seek = seek_rx.recv() => {
735 if let Some(seek) = seek {
736 apply_shard_seek(
737 &client, &stream, &shard, seek, &mut iterator, &mut tracker,
738 )
739 .await;
740 }
741 }
742 () = tokio::time::sleep(descriptor.poll_value()) => {}
743 }
744 }
745 }
746 Err(err) => {
747 let expired = err
748 .as_service_error()
749 .is_some_and(GetRecordsError::is_expired_iterator_exception);
750 let fatal = err
751 .as_service_error()
752 .is_some_and(GetRecordsError::is_resource_not_found_exception);
753 if !expired
754 && out
755 .send(Stamped::unstamped(Err(KinesisError::Read {
756 stream: stream.clone(),
757 shard: shard.clone(),
758 source: sdk_err(&err),
759 })))
760 .await
761 .is_err()
762 {
763 return;
764 }
765 if fatal {
766 let _ = store.release(&shard, &owner).await;
767 return;
768 }
769 failures += 1;
770 if failures > 10 {
771 let _ = store.release(&shard, &owner).await;
772 return;
773 }
774 tokio::time::sleep(Duration::from_secs(1)).await;
777 match initial_iterator(
778 &client,
779 &stream,
780 &shard,
781 &bus,
782 store.as_ref(),
783 Reopen::Recover,
784 )
785 .await
786 {
787 Ok(Some(fresh)) => iterator = fresh,
788 Ok(None) => return,
789 Err(_) => {}
790 }
791 }
792 }
793 }
794}