1use std::collections::BTreeSet;
5use std::ops::Bound;
6use std::sync::Arc;
7use std::time::Duration;
8
9use anyhow::Context;
10use anyhow::bail;
11use anyhow::ensure;
12use cohort::cohorts;
13use ingestion::ArcStreamingClient;
14use ingestion::ClientArgs;
15use ingestion::IngestionConfig;
16use ingestion::IngestionFactory;
17use ingestion::IngestionService;
18use ingestion::ingestion_client::IngestionClient;
19use metrics::IndexerMetrics;
20use prometheus::Registry;
21use sui_indexer_alt_framework_store_traits::ConcurrentStore;
22use sui_indexer_alt_framework_store_traits::Connection;
23use sui_indexer_alt_framework_store_traits::SequentialStore;
24use sui_indexer_alt_framework_store_traits::Store;
25use sui_indexer_alt_framework_store_traits::pipeline_task;
26use tracing::info;
27
28use crate::metrics::IngestionMetrics;
29use crate::pipeline::Processor;
30use crate::pipeline::concurrent::ConcurrentConfig;
31use crate::pipeline::concurrent::{self};
32use crate::pipeline::sequential::SequentialConfig;
33use crate::pipeline::sequential::{self};
34use crate::service::Service;
35
36pub use sui_field_count::FieldCount;
37pub use sui_futures::service;
38pub use sui_indexer_alt_framework_store_traits as store;
40pub use sui_types as types;
41
42#[cfg(feature = "cluster")]
43pub mod cluster;
44mod cohort;
45pub mod config;
46pub mod ingestion;
47pub mod metrics;
48pub mod pipeline;
49#[cfg(feature = "postgres")]
50pub mod postgres;
51
52#[cfg(test)]
53pub mod mocks;
54
55#[derive(clap::Args, Default, Debug, Clone)]
57pub struct IndexerArgs {
58 #[arg(long)]
65 pub first_checkpoint: Option<u64>,
66
67 #[arg(long)]
70 pub last_checkpoint: Option<u64>,
71
72 #[arg(long, action = clap::ArgAction::Append)]
75 pub pipeline: Vec<String>,
76
77 #[clap(flatten)]
79 pub task: TaskArgs,
80}
81
82#[derive(clap::Parser, Default, Debug, Clone)]
84pub struct TaskArgs {
85 #[arg(long, requires = "reader_interval_ms")]
97 task: Option<String>,
98
99 #[arg(long, requires = "task")]
110 reader_interval_ms: Option<u64>,
111}
112
113pub struct Indexer<S: Store> {
114 store: S,
118
119 metrics: Arc<IndexerMetrics>,
121
122 ingestion_factory: IngestionFactory,
125
126 first_checkpoint: Option<u64>,
129
130 last_checkpoint: Option<u64>,
133
134 latest_checkpoint: u64,
136
137 next_checkpoint: u64,
140
141 next_sequential_checkpoint: Option<u64>,
144
145 task: Option<Task>,
157
158 enabled_pipelines: Option<BTreeSet<String>>,
162
163 added_pipelines: BTreeSet<&'static str>,
166
167 pipelines: Vec<PendingPipeline>,
170}
171
172#[derive(Clone)]
174pub(crate) struct Task {
175 task: String,
178 reader_interval: Duration,
181}
182
183struct PendingPipeline {
187 name: &'static str,
189
190 next_checkpoint: u64,
192
193 build: PipelineBuilder,
196}
197
198type PipelineBuilder = Box<dyn FnOnce(&mut IngestionService) -> Service + Send + Sync>;
206
207impl TaskArgs {
208 pub fn tasked(task: String, reader_interval_ms: u64) -> Self {
209 Self {
210 task: Some(task),
211 reader_interval_ms: Some(reader_interval_ms),
212 }
213 }
214
215 fn into_task(self) -> Option<Task> {
216 Some(Task {
217 task: self.task?,
218 reader_interval: Duration::from_millis(self.reader_interval_ms?),
219 })
220 }
221}
222
223impl<S: Store> Indexer<S> {
224 pub async fn new(
236 store: S,
237 indexer_args: IndexerArgs,
238 client_args: ClientArgs,
239 ingestion_config: IngestionConfig,
240 metrics_prefix: Option<&str>,
241 registry: &Registry,
242 ) -> anyhow::Result<Self> {
243 let ingestion_factory =
244 IngestionFactory::new(client_args, ingestion_config, metrics_prefix, registry).await?;
245 Self::with_ingestion_factory(
246 store,
247 indexer_args,
248 ingestion_factory,
249 metrics_prefix,
250 registry,
251 )
252 .await
253 }
254
255 pub async fn with_ingestion_clients(
269 store: S,
270 indexer_args: IndexerArgs,
271 ingestion_client: IngestionClient,
272 streaming_client: Option<ArcStreamingClient>,
273 gap_fill_client: Option<IngestionClient>,
274 ingestion_config: IngestionConfig,
275 metrics_prefix: Option<&str>,
276 registry: &Registry,
277 ) -> anyhow::Result<Self> {
278 let ingestion_factory = IngestionFactory::with_clients(
279 ingestion_client,
280 streaming_client,
281 gap_fill_client,
282 ingestion_config,
283 );
284 Self::with_ingestion_factory(
285 store,
286 indexer_args,
287 ingestion_factory,
288 metrics_prefix,
289 registry,
290 )
291 .await
292 }
293
294 async fn with_ingestion_factory(
297 store: S,
298 indexer_args: IndexerArgs,
299 ingestion_factory: IngestionFactory,
300 metrics_prefix: Option<&str>,
301 registry: &Registry,
302 ) -> anyhow::Result<Self> {
303 let IndexerArgs {
304 first_checkpoint,
305 last_checkpoint,
306 pipeline,
307 task,
308 } = indexer_args;
309
310 let metrics = IndexerMetrics::new(metrics_prefix, registry);
311
312 let latest_checkpoint = ingestion_factory.latest_checkpoint_number().await?;
313
314 info!(latest_checkpoint);
315
316 Ok(Self {
317 store,
318 metrics,
319 ingestion_factory,
320 first_checkpoint,
321 last_checkpoint,
322 latest_checkpoint,
323 next_checkpoint: u64::MAX,
324 next_sequential_checkpoint: None,
325 task: task.into_task(),
326 enabled_pipelines: if pipeline.is_empty() {
327 None
328 } else {
329 Some(pipeline.into_iter().collect())
330 },
331 added_pipelines: BTreeSet::new(),
332 pipelines: vec![],
333 })
334 }
335
336 pub fn store(&self) -> &S {
338 &self.store
339 }
340
341 pub fn ingestion_client(&self) -> &IngestionClient {
343 self.ingestion_factory.ingestion_client()
344 }
345
346 pub fn indexer_metrics(&self) -> &Arc<IndexerMetrics> {
348 &self.metrics
349 }
350
351 pub fn ingestion_metrics(&self) -> &Arc<IngestionMetrics> {
353 self.ingestion_factory.metrics()
354 }
355
356 pub fn pipelines(&self) -> impl Iterator<Item = &'static str> + '_ {
358 self.added_pipelines.iter().copied().filter(|p| {
359 self.enabled_pipelines
360 .as_ref()
361 .is_none_or(|e| e.contains(*p))
362 })
363 }
364
365 pub fn next_sequential_checkpoint(&self) -> Option<u64> {
369 self.next_sequential_checkpoint
370 }
371
372 pub async fn run(self) -> anyhow::Result<Service> {
382 let Self {
383 ingestion_factory,
384 last_checkpoint,
385 latest_checkpoint,
386 enabled_pipelines,
387 pipelines,
388 ..
389 } = self;
390
391 if let Some(enabled_pipelines) = enabled_pipelines {
392 ensure!(
393 enabled_pipelines.is_empty(),
394 "Tried to enable pipelines that this indexer does not know about: \
395 {enabled_pipelines:#?}",
396 );
397 }
398
399 ensure!(!pipelines.is_empty(), "No pipelines registered to run");
400
401 let groups = cohorts(
402 pipelines,
403 latest_checkpoint,
404 ingestion_factory.config().min_cohort_boundary,
405 );
406
407 info!(
408 cohorts = groups.len(),
409 latest_checkpoint, "Grouped pipelines into ingestion cohorts"
410 );
411
412 let end = last_checkpoint.map_or(Bound::Unbounded, Bound::Included);
413
414 let mut service = Service::new();
415 for (cohort, group) in groups.into_iter().enumerate() {
416 let mut ingestion = ingestion_factory.create(cohort);
417
418 let mut start = u64::MAX;
419 let mut names = Vec::with_capacity(group.len());
420 let mut members = Vec::with_capacity(group.len());
421 for pending in group {
422 start = start.min(pending.next_checkpoint);
423 names.push(pending.name);
424 members.push((pending.build)(&mut ingestion));
425 }
426
427 info!(cohort, start, end = last_checkpoint, pipelines = ?names, "Ingestion range");
428
429 let mut cohort_service = ingestion
430 .run((Bound::Included(start), end))
431 .await
432 .context("Failed to start ingestion service")?;
433
434 for member in members {
435 cohort_service = cohort_service.merge(member);
436 }
437
438 service = service.merge(cohort_service);
439 }
440
441 Ok(service)
442 }
443
444 async fn add_pipeline<P: Processor + 'static>(
453 &mut self,
454 pipeline_task: String,
455 retention: Option<u64>,
456 ) -> anyhow::Result<Option<u64>> {
457 ensure!(
458 self.added_pipelines.insert(P::NAME),
459 "Pipeline {:?} already added",
460 P::NAME,
461 );
462
463 if let Some(enabled_pipelines) = &mut self.enabled_pipelines
464 && !enabled_pipelines.remove(P::NAME)
465 {
466 info!(pipeline = P::NAME, "Skipping");
467 return Ok(None);
468 }
469
470 let proposed_next_checkpoint = if let Some(first_checkpoint) = self.first_checkpoint {
473 first_checkpoint
474 } else if let Some(retention) = retention {
475 self.latest_checkpoint.saturating_sub(retention)
476 } else {
477 0
478 };
479 let mut conn = self.store.connect().await?;
480 let init_watermark = conn
481 .init_watermark(&pipeline_task, proposed_next_checkpoint.checked_sub(1))
482 .await
483 .with_context(|| format!("Failed to init watermark for {pipeline_task}"))?;
484
485 let next_checkpoint = if let Some(init_watermark) = init_watermark {
486 if let Some(checkpoint_hi_inclusive) = init_watermark.checkpoint_hi_inclusive {
487 checkpoint_hi_inclusive + 1
488 } else {
489 0
490 }
491 } else {
492 proposed_next_checkpoint
493 };
494
495 self.next_checkpoint = next_checkpoint.min(self.next_checkpoint);
496
497 Ok(Some(next_checkpoint))
498 }
499}
500
501impl<S: ConcurrentStore> Indexer<S> {
502 pub async fn concurrent_pipeline<H: concurrent::Handler<Store = S>>(
509 &mut self,
510 handler: H,
511 config: ConcurrentConfig,
512 ) -> anyhow::Result<()> {
513 let pipeline_task =
514 pipeline_task::<S>(H::NAME, self.task.as_ref().map(|t| t.task.as_str()))?;
515 let retention = config.pruner.as_ref().map(|p| p.retention);
516 let Some(next_checkpoint) = self.add_pipeline::<H>(pipeline_task, retention).await? else {
517 return Ok(());
518 };
519
520 let store = self.store.clone();
521 let task = self.task.clone();
522 let metrics = self.metrics.clone();
523 self.pipelines.push(PendingPipeline {
524 name: H::NAME,
525 next_checkpoint,
526 build: Box::new(move |ingestion| {
527 let checkpoint_rx =
528 ingestion.subscribe_bounded(config.ingestion.subscriber_channel_size());
529 concurrent::pipeline::<H>(
530 handler,
531 next_checkpoint,
532 config,
533 store,
534 task,
535 checkpoint_rx,
536 metrics,
537 )
538 }),
539 });
540
541 Ok(())
542 }
543}
544
545impl<T: SequentialStore> Indexer<T> {
546 pub async fn sequential_pipeline<H: sequential::Handler<Store = T>>(
554 &mut self,
555 handler: H,
556 config: SequentialConfig,
557 ) -> anyhow::Result<()> {
558 if self.task.is_some() {
559 bail!(
560 "Sequential pipelines do not support pipeline tasks. \
561 These pipelines guarantee that each checkpoint is committed exactly once and in order. \
562 Running the same pipeline under a different task would violate these guarantees."
563 );
564 }
565
566 let Some(next_checkpoint) = self.add_pipeline::<H>(H::NAME.to_owned(), None).await? else {
567 return Ok(());
568 };
569
570 self.next_sequential_checkpoint = Some(
572 self.next_sequential_checkpoint
573 .map_or(next_checkpoint, |n| n.min(next_checkpoint)),
574 );
575
576 let store = self.store.clone();
577 let metrics = self.metrics.clone();
578 self.pipelines.push(PendingPipeline {
579 name: H::NAME,
580 next_checkpoint,
581 build: Box::new(move |ingestion| {
582 let checkpoint_rx =
583 ingestion.subscribe_bounded(config.ingestion.subscriber_channel_size());
584 sequential::pipeline::<H>(
585 handler,
586 next_checkpoint,
587 config,
588 store,
589 checkpoint_rx,
590 metrics,
591 )
592 }),
593 });
594
595 Ok(())
596 }
597}
598
599#[cfg(test)]
600mod tests {
601 use std::sync::Arc;
602
603 use async_trait::async_trait;
604 use clap::Parser;
605 use sui_indexer_alt_framework_store_traits::PrunerWatermark;
606 use sui_synthetic_ingestion::synthetic_ingestion;
607 use tokio::sync::watch;
608
609 use crate::FieldCount;
610 use crate::config::ConcurrencyConfig;
611 use crate::ingestion::ingestion_client::IngestionClientArgs;
612 use crate::ingestion::store_client::ObjectStoreWatermark;
613 use crate::ingestion::store_client::WATERMARK_PATH;
614 use crate::mocks::store::FallibleMockStore;
615 use crate::pipeline::CommitterConfig;
616 use crate::pipeline::Processor;
617 use crate::pipeline::concurrent::ConcurrentConfig;
618 use crate::store::CommitterWatermark;
619 use crate::store::ConcurrentConnection as _;
620
621 use super::*;
622
623 #[allow(dead_code)]
624 #[derive(Clone, FieldCount)]
625 struct MockValue(u64);
626
627 struct ControllableHandler {
629 process_below: watch::Receiver<u64>,
631 }
632
633 impl ControllableHandler {
634 fn with_limit(limit: u64) -> (Self, watch::Sender<u64>) {
635 let (tx, rx) = watch::channel(limit);
636 (Self { process_below: rx }, tx)
637 }
638 }
639
640 #[async_trait]
641 impl Processor for ControllableHandler {
642 const NAME: &'static str = "controllable";
643 type Value = MockValue;
644
645 async fn process(
646 &self,
647 checkpoint: &Arc<sui_types::full_checkpoint_content::Checkpoint>,
648 ) -> anyhow::Result<Vec<Self::Value>> {
649 let cp_num = checkpoint.summary.sequence_number;
650
651 self.process_below
653 .clone()
654 .wait_for(|&limit| cp_num <= limit)
655 .await
656 .ok();
657
658 Ok(vec![MockValue(cp_num)])
659 }
660 }
661
662 #[async_trait]
663 impl concurrent::Handler for ControllableHandler {
664 type Store = FallibleMockStore;
665 type Batch = Vec<MockValue>;
666
667 fn batch(
668 &self,
669 batch: &mut Self::Batch,
670 values: &mut std::vec::IntoIter<Self::Value>,
671 ) -> concurrent::BatchStatus {
672 batch.extend(values);
673 concurrent::BatchStatus::Ready
674 }
675
676 async fn commit<'a>(
677 &self,
678 batch: &Self::Batch,
679 conn: &mut <Self::Store as Store>::Connection<'a>,
680 ) -> anyhow::Result<usize> {
681 for value in batch {
682 conn.0
683 .commit_data(Self::NAME, value.0, vec![value.0])
684 .await?;
685 }
686 Ok(batch.len())
687 }
688 }
689
690 macro_rules! test_pipeline {
691 ($handler:ident, $name:literal) => {
692 struct $handler;
693
694 #[async_trait]
695 impl Processor for $handler {
696 const NAME: &'static str = $name;
697 type Value = MockValue;
698 async fn process(
699 &self,
700 checkpoint: &Arc<sui_types::full_checkpoint_content::Checkpoint>,
701 ) -> anyhow::Result<Vec<Self::Value>> {
702 Ok(vec![MockValue(checkpoint.summary.sequence_number)])
703 }
704 }
705
706 #[async_trait]
707 impl crate::pipeline::concurrent::Handler for $handler {
708 type Store = FallibleMockStore;
709 type Batch = Vec<Self::Value>;
710
711 fn batch(
712 &self,
713 batch: &mut Self::Batch,
714 values: &mut std::vec::IntoIter<Self::Value>,
715 ) -> crate::pipeline::concurrent::BatchStatus {
716 batch.extend(values);
717 crate::pipeline::concurrent::BatchStatus::Pending
718 }
719
720 async fn commit<'a>(
721 &self,
722 batch: &Self::Batch,
723 conn: &mut <Self::Store as Store>::Connection<'a>,
724 ) -> anyhow::Result<usize> {
725 for value in batch {
726 conn.0
727 .commit_data(Self::NAME, value.0, vec![value.0])
728 .await?;
729 }
730 Ok(batch.len())
731 }
732 }
733
734 #[async_trait]
735 impl crate::pipeline::sequential::Handler for $handler {
736 type Store = FallibleMockStore;
737 type Batch = Vec<Self::Value>;
738
739 fn batch(&self, batch: &mut Self::Batch, values: std::vec::IntoIter<Self::Value>) {
740 batch.extend(values);
741 }
742
743 async fn commit<'a>(
744 &self,
745 _batch: &Self::Batch,
746 _conn: &mut <Self::Store as Store>::Connection<'a>,
747 ) -> anyhow::Result<usize> {
748 Ok(1)
749 }
750 }
751 };
752 }
753
754 test_pipeline!(MockHandler, "test_processor");
755 test_pipeline!(SequentialHandler, "sequential_handler");
756 test_pipeline!(MockCheckpointSequenceNumberHandler, "test");
757
758 fn init_ingestion_dir(latest_checkpoint: Option<u64>) -> tempfile::TempDir {
759 let dir = tempfile::tempdir().unwrap();
760 if let Some(cp) = latest_checkpoint {
761 let watermark_path = dir.path().join(WATERMARK_PATH);
762 std::fs::create_dir_all(watermark_path.parent().unwrap()).unwrap();
763 let watermark = ObjectStoreWatermark {
764 checkpoint_hi_inclusive: cp,
765 };
766 std::fs::write(watermark_path, serde_json::to_string(&watermark).unwrap()).unwrap();
767 }
768 dir
769 }
770
771 async fn create_test_indexer(
774 store: FallibleMockStore,
775 indexer_args: IndexerArgs,
776 registry: &Registry,
777 ingestion_data: Option<(u64, u64)>,
778 ) -> (Indexer<FallibleMockStore>, tempfile::TempDir) {
779 let temp_dir = init_ingestion_dir(None);
780 if let Some((num_checkpoints, checkpoint_size)) = ingestion_data {
781 synthetic_ingestion::generate_ingestion(synthetic_ingestion::Config {
782 ingestion_dir: temp_dir.path().to_owned(),
783 starting_checkpoint: 0,
784 num_checkpoints,
785 checkpoint_size,
786 })
787 .await;
788 }
789 let client_args = ClientArgs {
790 ingestion: IngestionClientArgs {
791 local_ingestion_path: Some(temp_dir.path().to_owned()),
792 ..Default::default()
793 },
794 ..Default::default()
795 };
796 let indexer = Indexer::new(
797 store,
798 indexer_args,
799 client_args,
800 IngestionConfig::default(),
801 None,
802 registry,
803 )
804 .await
805 .unwrap();
806 (indexer, temp_dir)
807 }
808
809 async fn set_committer_watermark(
810 conn: &mut <FallibleMockStore as Store>::Connection<'_>,
811 name: &str,
812 hi: u64,
813 ) {
814 conn.set_committer_watermark(
815 name,
816 CommitterWatermark {
817 checkpoint_hi_inclusive: hi,
818 ..Default::default()
819 },
820 )
821 .await
822 .unwrap();
823 }
824
825 async fn add_concurrent<H: concurrent::Handler<Store = FallibleMockStore>>(
826 indexer: &mut Indexer<FallibleMockStore>,
827 handler: H,
828 ) {
829 indexer
830 .concurrent_pipeline(handler, ConcurrentConfig::default())
831 .await
832 .unwrap();
833 }
834
835 async fn add_sequential<H: sequential::Handler<Store = FallibleMockStore>>(
836 indexer: &mut Indexer<FallibleMockStore>,
837 handler: H,
838 ) {
839 indexer
840 .sequential_pipeline(handler, SequentialConfig::default())
841 .await
842 .unwrap();
843 }
844
845 async fn multi_cohort_setup(
852 store: &FallibleMockStore,
853 near: &str,
854 far: &str,
855 ) -> tempfile::TempDir {
856 let mut conn = store.connect().await.unwrap();
857 set_committer_watermark(&mut conn, near, 100_100).await;
858 set_committer_watermark(&mut conn, far, 9).await;
859
860 let temp_dir = init_ingestion_dir(Some(100_000));
861 synthetic_ingestion::generate_ingestion(synthetic_ingestion::Config {
862 ingestion_dir: temp_dir.path().to_owned(),
863 starting_checkpoint: 0,
864 num_checkpoints: 60,
865 checkpoint_size: 1,
866 })
867 .await;
868
869 temp_dir
870 }
871
872 macro_rules! assert_out_of_order {
873 ($metrics:expr, $pipeline:expr, $expected:expr) => {
874 assert_eq!(
875 $metrics
876 .total_watermarks_out_of_order
877 .get_metric_with_label_values(&[$pipeline])
878 .unwrap()
879 .get(),
880 $expected,
881 );
882 };
883 }
884
885 async fn test_init_watermark(
886 first_checkpoint: Option<u64>,
887 is_concurrent: bool,
888 ) -> (Option<CommitterWatermark>, Option<PrunerWatermark>) {
889 let registry = Registry::new();
890 let store = FallibleMockStore::default();
891
892 test_pipeline!(A, "pipeline_name");
893
894 let mut conn = store.connect().await.unwrap();
895
896 let indexer_args = IndexerArgs {
897 first_checkpoint,
898 ..IndexerArgs::default()
899 };
900 let (mut indexer, _temp_dir) =
901 create_test_indexer(store.clone(), indexer_args, ®istry, None).await;
902
903 if is_concurrent {
904 add_concurrent(&mut indexer, A).await;
905 } else {
906 add_sequential(&mut indexer, A).await;
907 }
908
909 (
910 conn.committer_watermark(A::NAME).await.unwrap(),
911 conn.pruner_watermark(A::NAME, Duration::ZERO)
912 .await
913 .unwrap(),
914 )
915 }
916
917 const LATEST_CHECKPOINT: u64 = 10;
918
919 async fn test_next_checkpoint(
924 watermark: Option<u64>,
925 first_checkpoint: Option<u64>,
926 concurrent_config: ConcurrentConfig,
927 ) -> Indexer<FallibleMockStore> {
928 let registry = Registry::new();
929 let store = FallibleMockStore::default();
930
931 test_pipeline!(A, "concurrent_a");
932
933 if let Some(checkpoint_hi_inclusive) = watermark {
934 let mut conn = store.connect().await.unwrap();
935 conn.set_committer_watermark(
936 A::NAME,
937 CommitterWatermark {
938 checkpoint_hi_inclusive,
939 ..Default::default()
940 },
941 )
942 .await
943 .unwrap();
944 }
945
946 let temp_dir = init_ingestion_dir(Some(LATEST_CHECKPOINT));
947 let mut indexer = Indexer::new(
948 store,
949 IndexerArgs {
950 first_checkpoint,
951 ..Default::default()
952 },
953 ClientArgs {
954 ingestion: IngestionClientArgs {
955 local_ingestion_path: Some(temp_dir.path().to_owned()),
956 ..Default::default()
957 },
958 ..Default::default()
959 },
960 IngestionConfig::default(),
961 None,
962 ®istry,
963 )
964 .await
965 .unwrap();
966
967 assert_eq!(indexer.latest_checkpoint, LATEST_CHECKPOINT);
968
969 indexer
970 .concurrent_pipeline::<A>(A, concurrent_config)
971 .await
972 .unwrap();
973
974 indexer
975 }
976
977 fn pruner_config(retention: u64) -> ConcurrentConfig {
978 ConcurrentConfig {
979 pruner: Some(concurrent::PrunerConfig {
980 retention,
981 ..Default::default()
982 }),
983 ..Default::default()
984 }
985 }
986
987 #[test]
988 fn test_arg_parsing() {
989 #[derive(Parser)]
990 struct Args {
991 #[clap(flatten)]
992 indexer: IndexerArgs,
993 }
994
995 let args = Args::try_parse_from([
996 "cmd",
997 "--first-checkpoint",
998 "10",
999 "--last-checkpoint",
1000 "100",
1001 "--pipeline",
1002 "a",
1003 "--pipeline",
1004 "b",
1005 "--task",
1006 "t",
1007 "--reader-interval-ms",
1008 "5000",
1009 ])
1010 .unwrap();
1011
1012 assert_eq!(args.indexer.first_checkpoint, Some(10));
1013 assert_eq!(args.indexer.last_checkpoint, Some(100));
1014 assert_eq!(args.indexer.pipeline, vec!["a", "b"]);
1015 assert_eq!(args.indexer.task.task, Some("t".to_owned()));
1016 assert_eq!(args.indexer.task.reader_interval_ms, Some(5000));
1017 }
1018
1019 #[tokio::test]
1021 async fn test_next_checkpoint_all_pipelines_have_watermarks() {
1022 let registry = Registry::new();
1023 let store = FallibleMockStore::default();
1024
1025 test_pipeline!(A, "concurrent_a");
1026 test_pipeline!(B, "concurrent_b");
1027 test_pipeline!(C, "sequential_c");
1028 test_pipeline!(D, "sequential_d");
1029
1030 let mut conn = store.connect().await.unwrap();
1031
1032 conn.init_watermark(A::NAME, Some(0)).await.unwrap();
1033 set_committer_watermark(&mut conn, A::NAME, 100).await;
1034
1035 conn.init_watermark(B::NAME, Some(0)).await.unwrap();
1036 set_committer_watermark(&mut conn, B::NAME, 10).await;
1037
1038 conn.init_watermark(C::NAME, Some(0)).await.unwrap();
1039 set_committer_watermark(&mut conn, C::NAME, 1).await;
1040
1041 conn.init_watermark(D::NAME, Some(0)).await.unwrap();
1042 set_committer_watermark(&mut conn, D::NAME, 50).await;
1043
1044 let (mut indexer, _temp_dir) =
1045 create_test_indexer(store, IndexerArgs::default(), ®istry, None).await;
1046
1047 add_concurrent(&mut indexer, A).await;
1048 add_concurrent(&mut indexer, B).await;
1049 add_sequential(&mut indexer, C).await;
1050 add_sequential(&mut indexer, D).await;
1051
1052 assert_eq!(indexer.first_checkpoint, None);
1053 assert_eq!(indexer.last_checkpoint, None);
1054 assert_eq!(indexer.latest_checkpoint, 0);
1055 assert_eq!(indexer.next_checkpoint, 2);
1056 assert_eq!(indexer.next_sequential_checkpoint, Some(2));
1057 }
1058
1059 #[tokio::test]
1061 async fn test_next_checkpoint_not_all_pipelines_have_watermarks() {
1062 let registry = Registry::new();
1063 let store = FallibleMockStore::default();
1064
1065 test_pipeline!(A, "concurrent_a");
1066 test_pipeline!(B, "concurrent_b");
1067 test_pipeline!(C, "sequential_c");
1068 test_pipeline!(D, "sequential_d");
1069
1070 let mut conn = store.connect().await.unwrap();
1071 set_committer_watermark(&mut conn, B::NAME, 10).await;
1072 set_committer_watermark(&mut conn, C::NAME, 1).await;
1073
1074 let (mut indexer, _temp_dir) =
1075 create_test_indexer(store, IndexerArgs::default(), ®istry, None).await;
1076
1077 add_concurrent(&mut indexer, A).await;
1078 add_concurrent(&mut indexer, B).await;
1079 add_sequential(&mut indexer, C).await;
1080 add_sequential(&mut indexer, D).await;
1081
1082 assert_eq!(indexer.first_checkpoint, None);
1083 assert_eq!(indexer.last_checkpoint, None);
1084 assert_eq!(indexer.latest_checkpoint, 0);
1085 assert_eq!(indexer.next_checkpoint, 0);
1086 assert_eq!(indexer.next_sequential_checkpoint, Some(0));
1087 }
1088
1089 #[tokio::test]
1091 async fn test_next_checkpoint_smallest_is_0() {
1092 let registry = Registry::new();
1093 let store = FallibleMockStore::default();
1094
1095 test_pipeline!(A, "concurrent_a");
1096 test_pipeline!(B, "concurrent_b");
1097 test_pipeline!(C, "sequential_c");
1098 test_pipeline!(D, "sequential_d");
1099
1100 let mut conn = store.connect().await.unwrap();
1101 set_committer_watermark(&mut conn, A::NAME, 100).await;
1102 set_committer_watermark(&mut conn, B::NAME, 10).await;
1103 set_committer_watermark(&mut conn, C::NAME, 1).await;
1104 set_committer_watermark(&mut conn, D::NAME, 0).await;
1105
1106 let (mut indexer, _temp_dir) =
1107 create_test_indexer(store, IndexerArgs::default(), ®istry, None).await;
1108
1109 add_concurrent(&mut indexer, A).await;
1110 add_concurrent(&mut indexer, B).await;
1111 add_sequential(&mut indexer, C).await;
1112 add_sequential(&mut indexer, D).await;
1113
1114 assert_eq!(indexer.next_checkpoint, 1);
1115 }
1116
1117 #[tokio::test]
1120 async fn test_next_checkpoint_first_checkpoint_and_no_watermark() {
1121 let registry = Registry::new();
1122 let store = FallibleMockStore::default();
1123
1124 test_pipeline!(A, "concurrent_a");
1125 test_pipeline!(B, "concurrent_b");
1126 test_pipeline!(C, "sequential_c");
1127 test_pipeline!(D, "sequential_d");
1128
1129 let mut conn = store.connect().await.unwrap();
1130 set_committer_watermark(&mut conn, B::NAME, 50).await;
1131 set_committer_watermark(&mut conn, C::NAME, 10).await;
1132
1133 let indexer_args = IndexerArgs {
1134 first_checkpoint: Some(5),
1135 ..Default::default()
1136 };
1137 let (mut indexer, _temp_dir) =
1138 create_test_indexer(store, indexer_args, ®istry, None).await;
1139
1140 add_concurrent(&mut indexer, A).await;
1141 add_concurrent(&mut indexer, B).await;
1142 add_sequential(&mut indexer, C).await;
1143 add_sequential(&mut indexer, D).await;
1144
1145 assert_eq!(indexer.first_checkpoint, Some(5));
1146 assert_eq!(indexer.last_checkpoint, None);
1147 assert_eq!(indexer.latest_checkpoint, 0);
1148 assert_eq!(indexer.next_checkpoint, 5);
1149 assert_eq!(indexer.next_sequential_checkpoint, Some(5));
1150 }
1151
1152 #[tokio::test]
1155 async fn test_next_checkpoint_ignore_first_checkpoint() {
1156 let registry = Registry::new();
1157 let store = FallibleMockStore::default();
1158
1159 test_pipeline!(B, "concurrent_b");
1160 test_pipeline!(C, "sequential_c");
1161
1162 let mut conn = store.connect().await.unwrap();
1163 set_committer_watermark(&mut conn, B::NAME, 50).await;
1164 set_committer_watermark(&mut conn, C::NAME, 10).await;
1165
1166 let indexer_args = IndexerArgs {
1167 first_checkpoint: Some(5),
1168 ..Default::default()
1169 };
1170 let (mut indexer, _temp_dir) =
1171 create_test_indexer(store, indexer_args, ®istry, None).await;
1172
1173 add_concurrent(&mut indexer, B).await;
1174 add_sequential(&mut indexer, C).await;
1175
1176 assert_eq!(indexer.first_checkpoint, Some(5));
1177 assert_eq!(indexer.last_checkpoint, None);
1178 assert_eq!(indexer.latest_checkpoint, 0);
1179 assert_eq!(indexer.next_checkpoint, 11);
1180 assert_eq!(indexer.next_sequential_checkpoint, Some(11));
1181 }
1182
1183 #[tokio::test]
1187 async fn test_next_checkpoint_large_first_checkpoint() {
1188 let registry = Registry::new();
1189 let store = FallibleMockStore::default();
1190
1191 test_pipeline!(A, "concurrent_a");
1192 test_pipeline!(B, "concurrent_b");
1193 test_pipeline!(C, "sequential_c");
1194
1195 let mut conn = store.connect().await.unwrap();
1196 set_committer_watermark(&mut conn, B::NAME, 50).await;
1197 set_committer_watermark(&mut conn, C::NAME, 10).await;
1198
1199 let indexer_args = IndexerArgs {
1200 first_checkpoint: Some(24),
1201 ..Default::default()
1202 };
1203 let (mut indexer, _temp_dir) =
1204 create_test_indexer(store, indexer_args, ®istry, None).await;
1205
1206 add_concurrent(&mut indexer, A).await;
1207 add_concurrent(&mut indexer, B).await;
1208 add_sequential(&mut indexer, C).await;
1209
1210 assert_eq!(indexer.first_checkpoint, Some(24));
1211 assert_eq!(indexer.last_checkpoint, None);
1212 assert_eq!(indexer.latest_checkpoint, 0);
1213 assert_eq!(indexer.next_checkpoint, 11);
1214 assert_eq!(indexer.next_sequential_checkpoint, Some(11));
1215 }
1216
1217 #[tokio::test]
1219 async fn test_latest_checkpoint_from_watermark() {
1220 let registry = Registry::new();
1221 let store = FallibleMockStore::default();
1222 let temp_dir = init_ingestion_dir(Some(30));
1223 let indexer = Indexer::new(
1224 store,
1225 IndexerArgs::default(),
1226 ClientArgs {
1227 ingestion: IngestionClientArgs {
1228 local_ingestion_path: Some(temp_dir.path().to_owned()),
1229 ..Default::default()
1230 },
1231 ..Default::default()
1232 },
1233 IngestionConfig::default(),
1234 None,
1235 ®istry,
1236 )
1237 .await
1238 .unwrap();
1239
1240 assert_eq!(indexer.latest_checkpoint, 30);
1241 }
1242
1243 #[tokio::test]
1246 async fn test_next_checkpoint_with_pruner_uses_retention() {
1247 let retention = LATEST_CHECKPOINT - 1;
1248 let indexer = test_next_checkpoint(None, None, pruner_config(retention)).await;
1249 assert_eq!(indexer.next_checkpoint, LATEST_CHECKPOINT - retention);
1250 }
1251
1252 #[tokio::test]
1254 async fn test_next_checkpoint_without_pruner_falls_back_to_genesis() {
1255 let indexer = test_next_checkpoint(None, None, ConcurrentConfig::default()).await;
1256 assert_eq!(indexer.next_checkpoint, 0);
1257 }
1258
1259 #[tokio::test]
1261 async fn test_next_checkpoint_watermark_takes_priority_over_pruner() {
1262 let retention = LATEST_CHECKPOINT - 1;
1263 let indexer = test_next_checkpoint(Some(5), None, pruner_config(retention)).await;
1264 assert_eq!(indexer.next_checkpoint, 6);
1265 }
1266
1267 #[tokio::test]
1270 async fn test_next_checkpoint_first_checkpoint_takes_priority_over_pruner() {
1271 let retention = LATEST_CHECKPOINT - 1;
1272 let indexer = test_next_checkpoint(None, Some(2), pruner_config(retention)).await;
1273 assert_eq!(indexer.next_checkpoint, 2);
1274 }
1275
1276 #[tokio::test]
1278 async fn test_next_checkpoint_retention_exceeds_latest_checkpoint() {
1279 let retention = LATEST_CHECKPOINT + 1;
1280 let indexer = test_next_checkpoint(None, None, pruner_config(retention)).await;
1281 assert_eq!(indexer.next_checkpoint, 0);
1282 }
1283
1284 #[tokio::test]
1286 async fn test_indexer_ingestion_existing_watermarks_no_first_checkpoint() {
1287 let registry = Registry::new();
1288 let store = FallibleMockStore::default();
1289
1290 test_pipeline!(A, "concurrent_a");
1291 test_pipeline!(B, "concurrent_b");
1292 test_pipeline!(C, "sequential_c");
1293 test_pipeline!(D, "sequential_d");
1294
1295 let mut conn = store.connect().await.unwrap();
1296 set_committer_watermark(&mut conn, A::NAME, 5).await;
1297 set_committer_watermark(&mut conn, B::NAME, 10).await;
1298 set_committer_watermark(&mut conn, C::NAME, 15).await;
1299 set_committer_watermark(&mut conn, D::NAME, 20).await;
1300
1301 let indexer_args = IndexerArgs {
1302 last_checkpoint: Some(29),
1303 ..Default::default()
1304 };
1305 let (mut indexer, _temp_dir) =
1306 create_test_indexer(store.clone(), indexer_args, ®istry, Some((30, 1))).await;
1307
1308 add_concurrent(&mut indexer, A).await;
1309 add_concurrent(&mut indexer, B).await;
1310 add_sequential(&mut indexer, C).await;
1311 add_sequential(&mut indexer, D).await;
1312
1313 let ingestion_metrics = indexer.ingestion_metrics().clone();
1314 let indexer_metrics = indexer.indexer_metrics().clone();
1315
1316 indexer.run().await.unwrap().join().await.unwrap();
1317
1318 assert_eq!(
1319 ingestion_metrics
1320 .total_ingested_checkpoints
1321 .with_label_values(&["0"])
1322 .get(),
1323 24
1324 );
1325 assert_out_of_order!(indexer_metrics, A::NAME, 0);
1326 assert_out_of_order!(indexer_metrics, B::NAME, 5);
1327 assert_out_of_order!(indexer_metrics, C::NAME, 10);
1328 assert_out_of_order!(indexer_metrics, D::NAME, 15);
1329 }
1330
1331 #[tokio::test]
1333 async fn test_indexer_ingestion_existing_watermarks_ignore_first_checkpoint() {
1334 let registry = Registry::new();
1335 let store = FallibleMockStore::default();
1336
1337 test_pipeline!(A, "concurrent_a");
1338 test_pipeline!(B, "concurrent_b");
1339 test_pipeline!(C, "sequential_c");
1340 test_pipeline!(D, "sequential_d");
1341
1342 let mut conn = store.connect().await.unwrap();
1343 set_committer_watermark(&mut conn, A::NAME, 5).await;
1344 set_committer_watermark(&mut conn, B::NAME, 10).await;
1345 set_committer_watermark(&mut conn, C::NAME, 15).await;
1346 set_committer_watermark(&mut conn, D::NAME, 20).await;
1347
1348 let indexer_args = IndexerArgs {
1349 first_checkpoint: Some(3),
1350 last_checkpoint: Some(29),
1351 ..Default::default()
1352 };
1353 let (mut indexer, _temp_dir) =
1354 create_test_indexer(store.clone(), indexer_args, ®istry, Some((30, 1))).await;
1355
1356 add_concurrent(&mut indexer, A).await;
1357 add_concurrent(&mut indexer, B).await;
1358 add_sequential(&mut indexer, C).await;
1359 add_sequential(&mut indexer, D).await;
1360
1361 let ingestion_metrics = indexer.ingestion_metrics().clone();
1362 let metrics = indexer.indexer_metrics().clone();
1363 indexer.run().await.unwrap().join().await.unwrap();
1364
1365 assert_eq!(
1366 ingestion_metrics
1367 .total_ingested_checkpoints
1368 .with_label_values(&["0"])
1369 .get(),
1370 24
1371 );
1372 assert_out_of_order!(metrics, A::NAME, 0);
1373 assert_out_of_order!(metrics, B::NAME, 5);
1374 assert_out_of_order!(metrics, C::NAME, 10);
1375 assert_out_of_order!(metrics, D::NAME, 15);
1376 }
1377
1378 #[tokio::test]
1380 async fn test_indexer_ingestion_missing_watermarks_no_first_checkpoint() {
1381 let registry = Registry::new();
1382 let store = FallibleMockStore::default();
1383
1384 test_pipeline!(A, "concurrent_a");
1385 test_pipeline!(B, "concurrent_b");
1386 test_pipeline!(C, "sequential_c");
1387 test_pipeline!(D, "sequential_d");
1388
1389 let mut conn = store.connect().await.unwrap();
1390 set_committer_watermark(&mut conn, B::NAME, 10).await;
1391 set_committer_watermark(&mut conn, C::NAME, 15).await;
1392 set_committer_watermark(&mut conn, D::NAME, 20).await;
1393
1394 let indexer_args = IndexerArgs {
1395 last_checkpoint: Some(29),
1396 ..Default::default()
1397 };
1398 let (mut indexer, _temp_dir) =
1399 create_test_indexer(store.clone(), indexer_args, ®istry, Some((30, 1))).await;
1400
1401 add_concurrent(&mut indexer, A).await;
1402 add_concurrent(&mut indexer, B).await;
1403 add_sequential(&mut indexer, C).await;
1404 add_sequential(&mut indexer, D).await;
1405
1406 let ingestion_metrics = indexer.ingestion_metrics().clone();
1407 let metrics = indexer.indexer_metrics().clone();
1408 indexer.run().await.unwrap().join().await.unwrap();
1409
1410 assert_eq!(
1411 ingestion_metrics
1412 .total_ingested_checkpoints
1413 .with_label_values(&["0"])
1414 .get(),
1415 30
1416 );
1417 assert_out_of_order!(metrics, A::NAME, 0);
1418 assert_out_of_order!(metrics, B::NAME, 11);
1419 assert_out_of_order!(metrics, C::NAME, 16);
1420 assert_out_of_order!(metrics, D::NAME, 21);
1421 }
1422
1423 #[tokio::test]
1425 async fn test_indexer_ingestion_use_first_checkpoint() {
1426 let registry = Registry::new();
1427 let store = FallibleMockStore::default();
1428
1429 test_pipeline!(A, "concurrent_a");
1430 test_pipeline!(B, "concurrent_b");
1431 test_pipeline!(C, "sequential_c");
1432 test_pipeline!(D, "sequential_d");
1433
1434 let mut conn = store.connect().await.unwrap();
1435 set_committer_watermark(&mut conn, B::NAME, 10).await;
1436 set_committer_watermark(&mut conn, C::NAME, 15).await;
1437 set_committer_watermark(&mut conn, D::NAME, 20).await;
1438
1439 let indexer_args = IndexerArgs {
1440 first_checkpoint: Some(10),
1441 last_checkpoint: Some(29),
1442 ..Default::default()
1443 };
1444 let (mut indexer, _temp_dir) =
1445 create_test_indexer(store.clone(), indexer_args, ®istry, Some((30, 1))).await;
1446
1447 add_concurrent(&mut indexer, A).await;
1448 add_concurrent(&mut indexer, B).await;
1449 add_sequential(&mut indexer, C).await;
1450 add_sequential(&mut indexer, D).await;
1451
1452 let ingestion_metrics = indexer.ingestion_metrics().clone();
1453 let metrics = indexer.indexer_metrics().clone();
1454 indexer.run().await.unwrap().join().await.unwrap();
1455
1456 assert_eq!(
1457 ingestion_metrics
1458 .total_ingested_checkpoints
1459 .with_label_values(&["0"])
1460 .get(),
1461 20
1462 );
1463 assert_out_of_order!(metrics, A::NAME, 0);
1464 assert_out_of_order!(metrics, B::NAME, 1);
1465 assert_out_of_order!(metrics, C::NAME, 6);
1466 assert_out_of_order!(metrics, D::NAME, 11);
1467 }
1468
1469 #[tokio::test]
1470 async fn test_init_watermark_concurrent_no_first_checkpoint() {
1471 let (committer_watermark, pruner_watermark) = test_init_watermark(None, true).await;
1472 assert_eq!(committer_watermark, None);
1473 assert_eq!(pruner_watermark, None);
1474 }
1475
1476 #[tokio::test]
1477 async fn test_init_watermark_concurrent_first_checkpoint_0() {
1478 let (committer_watermark, pruner_watermark) = test_init_watermark(Some(0), true).await;
1479 assert_eq!(committer_watermark, None);
1480 assert_eq!(pruner_watermark, None);
1481 }
1482
1483 #[tokio::test]
1484 async fn test_init_watermark_concurrent_first_checkpoint_1() {
1485 let (committer_watermark, pruner_watermark) = test_init_watermark(Some(1), true).await;
1486
1487 let committer_watermark = committer_watermark.unwrap();
1488 assert_eq!(committer_watermark.checkpoint_hi_inclusive, 0);
1489
1490 let pruner_watermark = pruner_watermark.unwrap();
1491 assert_eq!(pruner_watermark.reader_lo, 1);
1492 assert_eq!(pruner_watermark.pruner_hi, 1);
1493 }
1494
1495 #[tokio::test]
1496 async fn test_init_watermark_sequential() {
1497 let (committer_watermark, pruner_watermark) = test_init_watermark(Some(1), false).await;
1498
1499 let committer_watermark = committer_watermark.unwrap();
1500 assert_eq!(committer_watermark.checkpoint_hi_inclusive, 0);
1501
1502 let pruner_watermark = pruner_watermark.unwrap();
1503 assert_eq!(pruner_watermark.reader_lo, 1);
1504 assert_eq!(pruner_watermark.pruner_hi, 1);
1505 }
1506
1507 #[tokio::test]
1508 async fn test_multiple_sequential_pipelines_next_checkpoint() {
1509 let registry = Registry::new();
1510 let store = FallibleMockStore::default();
1511
1512 let mut conn = store.connect().await.unwrap();
1513 set_committer_watermark(&mut conn, MockHandler::NAME, 10).await;
1514 set_committer_watermark(&mut conn, SequentialHandler::NAME, 5).await;
1515
1516 let indexer_args = IndexerArgs {
1517 first_checkpoint: None,
1518 last_checkpoint: Some(19),
1519 pipeline: vec![],
1520 ..Default::default()
1521 };
1522 let (mut indexer, _temp_dir) =
1523 create_test_indexer(store.clone(), indexer_args, ®istry, Some((20, 2))).await;
1524
1525 add_sequential(&mut indexer, MockHandler).await;
1527
1528 assert_eq!(
1530 indexer.next_sequential_checkpoint(),
1531 Some(11),
1532 "next_sequential_checkpoint should be 11"
1533 );
1534
1535 add_sequential(&mut indexer, SequentialHandler).await;
1537
1538 assert_eq!(
1540 indexer.next_sequential_checkpoint(),
1541 Some(6),
1542 "next_sequential_checkpoint should still be 6"
1543 );
1544
1545 indexer.run().await.unwrap().join().await.unwrap();
1547
1548 let watermark1 = conn.committer_watermark(MockHandler::NAME).await.unwrap();
1550 let watermark2 = conn
1551 .committer_watermark(SequentialHandler::NAME)
1552 .await
1553 .unwrap();
1554
1555 assert_eq!(watermark1.unwrap().checkpoint_hi_inclusive, 19);
1556 assert_eq!(watermark2.unwrap().checkpoint_hi_inclusive, 19);
1557 }
1558
1559 #[tokio::test]
1563 async fn test_tasked_pipelines_ignore_below_main_reader_lo() {
1564 let registry = Registry::new();
1565 let store = FallibleMockStore::default();
1566
1567 let mut conn = store.connect().await.unwrap();
1570 set_committer_watermark(&mut conn, MockCheckpointSequenceNumberHandler::NAME, 10).await;
1571 conn.set_reader_watermark(MockCheckpointSequenceNumberHandler::NAME, 7)
1572 .await
1573 .unwrap();
1574
1575 let indexer_args = IndexerArgs {
1578 first_checkpoint: Some(0),
1579 last_checkpoint: Some(15),
1580 task: TaskArgs::tasked("task".to_string(), 10),
1581 ..Default::default()
1582 };
1583 let (mut tasked_indexer, _temp_dir) =
1584 create_test_indexer(store.clone(), indexer_args, ®istry, Some((16, 2))).await;
1585
1586 add_concurrent(&mut tasked_indexer, MockCheckpointSequenceNumberHandler).await;
1587
1588 let ingestion_metrics = tasked_indexer.ingestion_metrics().clone();
1589 let metrics = tasked_indexer.indexer_metrics().clone();
1590
1591 tasked_indexer.run().await.unwrap().join().await.unwrap();
1592
1593 assert_eq!(
1594 ingestion_metrics
1595 .total_ingested_checkpoints
1596 .with_label_values(&["0"])
1597 .get(),
1598 16
1599 );
1600 assert_eq!(
1601 metrics
1602 .total_collector_skipped_checkpoints
1603 .get_metric_with_label_values(&[MockCheckpointSequenceNumberHandler::NAME])
1604 .unwrap()
1605 .get(),
1606 7
1607 );
1608 let data = store
1609 .data
1610 .get(MockCheckpointSequenceNumberHandler::NAME)
1611 .unwrap();
1612 assert_eq!(data.len(), 9);
1613 for i in 0..7 {
1614 assert!(data.get(&i).is_none());
1615 }
1616 for i in 7..16 {
1617 assert!(data.get(&i).is_some());
1618 }
1619 }
1620
1621 #[tokio::test]
1623 async fn test_tasked_pipelines_surpass_main_pipeline_committer_hi() {
1624 let registry = Registry::new();
1625 let store = FallibleMockStore::default();
1626
1627 let mut conn = store.connect().await.unwrap();
1628 set_committer_watermark(&mut conn, "test", 10).await;
1629 conn.set_reader_watermark("test", 5).await.unwrap();
1630
1631 let indexer_args = IndexerArgs {
1634 first_checkpoint: Some(9),
1635 last_checkpoint: Some(25),
1636 task: TaskArgs::tasked("task".to_string(), 10),
1637 ..Default::default()
1638 };
1639 let (mut tasked_indexer, _temp_dir) =
1640 create_test_indexer(store.clone(), indexer_args, ®istry, Some((26, 2))).await;
1641
1642 add_concurrent(&mut tasked_indexer, MockCheckpointSequenceNumberHandler).await;
1643
1644 let ingestion_metrics = tasked_indexer.ingestion_metrics().clone();
1645 let metrics = tasked_indexer.indexer_metrics().clone();
1646
1647 tasked_indexer.run().await.unwrap().join().await.unwrap();
1648
1649 assert_eq!(
1650 ingestion_metrics
1651 .total_ingested_checkpoints
1652 .with_label_values(&["0"])
1653 .get(),
1654 17
1655 );
1656 assert_out_of_order!(metrics, "test", 0);
1657 assert_eq!(
1658 metrics
1659 .total_collector_skipped_checkpoints
1660 .get_metric_with_label_values(&[MockCheckpointSequenceNumberHandler::NAME])
1661 .unwrap()
1662 .get(),
1663 0
1664 );
1665
1666 let data = store.data.get("test").unwrap();
1667 assert_eq!(data.len(), 17);
1668 for i in 0..9 {
1669 assert!(data.get(&i).is_none());
1670 }
1671 for i in 9..26 {
1672 assert!(data.get(&i).is_some());
1673 }
1674 let main_pipeline_watermark = store.watermark("test").unwrap();
1675 assert_eq!(main_pipeline_watermark.checkpoint_hi_inclusive, Some(10));
1677 assert_eq!(main_pipeline_watermark.reader_lo, 5);
1678 let tasked_pipeline_watermark = store.watermark("test@task").unwrap();
1679 assert_eq!(tasked_pipeline_watermark.checkpoint_hi_inclusive, Some(25));
1680 assert_eq!(tasked_pipeline_watermark.reader_lo, 9);
1681 }
1682
1683 #[tokio::test]
1686 async fn test_tasked_pipelines_skip_checkpoints_trailing_main_reader_lo() {
1687 let registry = Registry::new();
1688 let store = FallibleMockStore::default();
1689 let mut conn = store.connect().await.unwrap();
1690 set_committer_watermark(&mut conn, ControllableHandler::NAME, 11).await;
1692
1693 let indexer_args = IndexerArgs {
1695 first_checkpoint: Some(0),
1696 last_checkpoint: Some(500),
1697 task: TaskArgs::tasked("task".to_string(), 10 ),
1698 ..Default::default()
1699 };
1700 let (mut tasked_indexer, _temp_dir) =
1701 create_test_indexer(store.clone(), indexer_args, ®istry, Some((501, 2))).await;
1702 let mut allow_process = 10;
1703 let (controllable_handler, process_below) = ControllableHandler::with_limit(allow_process);
1705 let _ = tasked_indexer
1706 .concurrent_pipeline(
1707 controllable_handler,
1708 ConcurrentConfig {
1709 committer: CommitterConfig {
1710 collect_interval_ms: 10,
1711 watermark_interval_ms: 10,
1712 ..Default::default()
1713 },
1714 fanout: Some(ConcurrencyConfig::Fixed { value: 501 }),
1717 ..Default::default()
1718 },
1719 )
1720 .await;
1721 let metrics = tasked_indexer.indexer_metrics().clone();
1722
1723 let mut s_indexer = tasked_indexer.run().await.unwrap();
1724
1725 store
1729 .wait_for_watermark(
1730 &pipeline_task::<FallibleMockStore>(ControllableHandler::NAME, Some("task"))
1731 .unwrap(),
1732 10,
1733 Duration::from_secs(10),
1734 )
1735 .await;
1736
1737 conn.set_reader_watermark(ControllableHandler::NAME, 250)
1742 .await
1743 .unwrap();
1744
1745 let reader_lo = metrics
1746 .collector_reader_lo
1747 .with_label_values(&[ControllableHandler::NAME]);
1748
1749 let mut interval = tokio::time::interval(Duration::from_millis(10));
1753 while reader_lo.get() != 250 {
1754 interval.tick().await;
1755 allow_process += 1;
1757 assert!(
1758 allow_process <= 500,
1759 "Released all checkpoints but collector never observed new reader_lo"
1760 );
1761 process_below.send(allow_process).ok();
1762 }
1763
1764 process_below.send(500).ok();
1771
1772 s_indexer.join().await.unwrap();
1773
1774 let data = store.data.get(ControllableHandler::NAME).unwrap();
1775
1776 for chkpt in (allow_process + 1)..250 {
1778 assert!(
1779 data.get(&chkpt).is_none(),
1780 "Checkpoint {chkpt} should have been skipped"
1781 );
1782 }
1783
1784 for chkpt in 250..=500 {
1786 assert!(
1787 data.get(&chkpt).is_some(),
1788 "Checkpoint {chkpt} should have been committed (>= reader_lo)"
1789 );
1790 }
1791
1792 for chkpt in 0..=10 {
1794 assert!(
1795 data.get(&chkpt).is_some(),
1796 "Checkpoint {chkpt} should have been committed (baseline)"
1797 );
1798 }
1799 }
1800
1801 #[tokio::test]
1805 async fn test_multi_cohort_ingestion() {
1806 let registry = Registry::new();
1807 let store = FallibleMockStore::default();
1808
1809 test_pipeline!(A, "near_tip");
1810 test_pipeline!(B, "far_behind");
1811
1812 let temp_dir = multi_cohort_setup(&store, A::NAME, B::NAME).await;
1813
1814 let mut indexer = Indexer::new(
1815 store.clone(),
1816 IndexerArgs {
1817 last_checkpoint: Some(59),
1818 ..Default::default()
1819 },
1820 ClientArgs {
1821 ingestion: IngestionClientArgs {
1822 local_ingestion_path: Some(temp_dir.path().to_owned()),
1823 ..Default::default()
1824 },
1825 ..Default::default()
1826 },
1827 IngestionConfig::default(),
1828 None,
1829 ®istry,
1830 )
1831 .await
1832 .unwrap();
1833
1834 add_concurrent(&mut indexer, A).await;
1835 add_concurrent(&mut indexer, B).await;
1836
1837 let ingestion_metrics = indexer.ingestion_metrics().clone();
1838 let indexer_metrics = indexer.indexer_metrics().clone();
1839 indexer.run().await.unwrap().join().await.unwrap();
1840
1841 assert_eq!(
1845 ingestion_metrics
1846 .total_ingested_checkpoints
1847 .with_label_values(&["0"])
1848 .get(),
1849 0
1850 );
1851 assert_eq!(
1852 ingestion_metrics
1853 .total_ingested_checkpoints
1854 .with_label_values(&["1"])
1855 .get(),
1856 50
1857 );
1858
1859 let received = |name| {
1862 indexer_metrics
1863 .total_handler_checkpoints_received
1864 .get_metric_with_label_values(&[name])
1865 .unwrap()
1866 .get()
1867 };
1868 assert_eq!(received(A::NAME), 0);
1869 assert_eq!(received(B::NAME), 50);
1870
1871 let watermark = store.watermark(B::NAME).unwrap();
1874 assert_eq!(watermark.checkpoint_hi_inclusive, Some(59));
1875 let watermark = store.watermark(A::NAME).unwrap();
1876 assert_eq!(watermark.checkpoint_hi_inclusive, Some(100_100));
1877 }
1878
1879 #[tokio::test]
1882 async fn test_multi_cohort_ingestion_with_clients() {
1883 let registry = Registry::new();
1884 let store = FallibleMockStore::default();
1885
1886 test_pipeline!(A, "near_tip");
1887 test_pipeline!(B, "far_behind");
1888
1889 let temp_dir = multi_cohort_setup(&store, A::NAME, B::NAME).await;
1890
1891 let ingestion_metrics = IngestionMetrics::new(None, ®istry);
1892 let ingestion_client = IngestionClient::new(
1893 IngestionClientArgs {
1894 local_ingestion_path: Some(temp_dir.path().to_owned()),
1895 ..Default::default()
1896 },
1897 ingestion_metrics.clone(),
1898 )
1899 .unwrap();
1900
1901 let mut indexer = Indexer::with_ingestion_clients(
1902 store.clone(),
1903 IndexerArgs {
1904 last_checkpoint: Some(59),
1905 ..Default::default()
1906 },
1907 ingestion_client,
1908 None,
1909 None,
1910 IngestionConfig::default(),
1911 None,
1912 ®istry,
1913 )
1914 .await
1915 .unwrap();
1916
1917 add_concurrent(&mut indexer, A).await;
1918 add_concurrent(&mut indexer, B).await;
1919
1920 let indexer_metrics = indexer.indexer_metrics().clone();
1921 indexer.run().await.unwrap().join().await.unwrap();
1922
1923 assert_eq!(
1926 ingestion_metrics
1927 .total_ingested_checkpoints
1928 .with_label_values(&["0"])
1929 .get(),
1930 0
1931 );
1932 assert_eq!(
1933 ingestion_metrics
1934 .total_ingested_checkpoints
1935 .with_label_values(&["1"])
1936 .get(),
1937 50
1938 );
1939 assert_eq!(
1940 indexer_metrics
1941 .total_handler_checkpoints_received
1942 .get_metric_with_label_values(&[A::NAME])
1943 .unwrap()
1944 .get(),
1945 0
1946 );
1947
1948 let watermark = store.watermark(B::NAME).unwrap();
1949 assert_eq!(watermark.checkpoint_hi_inclusive, Some(59));
1950 }
1951
1952 #[tokio::test]
1957 async fn test_near_tip_cohort_progresses_while_far_behind_stalls() {
1958 const TIP: u64 = 100;
1959
1960 let registry = Registry::new();
1961 let store = FallibleMockStore::default();
1962
1963 test_pipeline!(A, "near_tip");
1964
1965 let (far_behind, _release) = ControllableHandler::with_limit(9);
1969
1970 let temp_dir = init_ingestion_dir(Some(TIP));
1971 let mut conn = store.connect().await.unwrap();
1972 set_committer_watermark(&mut conn, A::NAME, 94).await; set_committer_watermark(&mut conn, ControllableHandler::NAME, 9).await; synthetic_ingestion::generate_ingestion(synthetic_ingestion::Config {
1975 ingestion_dir: temp_dir.path().to_owned(),
1976 starting_checkpoint: 0,
1977 num_checkpoints: TIP + 1, checkpoint_size: 1,
1979 })
1980 .await;
1981
1982 let mut indexer = Indexer::new(
1983 store.clone(),
1984 IndexerArgs {
1985 last_checkpoint: Some(TIP),
1986 ..Default::default()
1987 },
1988 ClientArgs {
1989 ingestion: IngestionClientArgs {
1990 local_ingestion_path: Some(temp_dir.path().to_owned()),
1991 ..Default::default()
1992 },
1993 ..Default::default()
1994 },
1995 IngestionConfig {
1998 min_cohort_boundary: 10,
1999 ..Default::default()
2000 },
2001 None,
2002 ®istry,
2003 )
2004 .await
2005 .unwrap();
2006
2007 indexer
2010 .concurrent_pipeline(
2011 A,
2012 ConcurrentConfig {
2013 committer: CommitterConfig {
2014 collect_interval_ms: 10,
2015 watermark_interval_ms: 10,
2016 ..Default::default()
2017 },
2018 ..Default::default()
2019 },
2020 )
2021 .await
2022 .unwrap();
2023 add_concurrent(&mut indexer, far_behind).await;
2025
2026 let service = indexer.run().await.unwrap();
2027
2028 store
2032 .wait_for_watermark(A::NAME, TIP, Duration::from_secs(10))
2033 .await;
2034
2035 assert_eq!(
2038 store
2039 .watermark(ControllableHandler::NAME)
2040 .unwrap()
2041 .checkpoint_hi_inclusive,
2042 Some(9),
2043 );
2044 assert!(store.data.get(ControllableHandler::NAME).is_none());
2045
2046 drop(service);
2048 }
2049}