sui_indexer_alt_framework/pipeline/concurrent/
mod.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5use std::time::Duration;
6
7use async_trait::async_trait;
8use serde::Deserialize;
9use serde::Serialize;
10use sui_futures::service::Service;
11use tokio::sync::SetOnce;
12use tokio::sync::mpsc;
13use tracing::info;
14
15use crate::Task;
16use crate::config::ConcurrencyConfig;
17use crate::ingestion::ingestion_client::CheckpointEnvelope;
18use crate::metrics::IndexerMetrics;
19use crate::pipeline::CommitterConfig;
20use crate::pipeline::IngestionConfig;
21use crate::pipeline::Processor;
22use crate::pipeline::WatermarkPart;
23use crate::pipeline::concurrent::collector::collector;
24use crate::pipeline::concurrent::commit_watermark::commit_watermark;
25use crate::pipeline::concurrent::committer::committer;
26use crate::pipeline::concurrent::main_reader_lo::track_main_reader_lo;
27use crate::pipeline::concurrent::pruner::pruner;
28use crate::pipeline::concurrent::reader_watermark::reader_watermark;
29use crate::pipeline::processor::processor;
30use crate::store::ConcurrentStore;
31use crate::store::Store;
32
33mod collector;
34mod commit_watermark;
35mod committer;
36mod main_reader_lo;
37mod pruner;
38mod reader_watermark;
39
40/// Status returned by `Handler::batch` to indicate whether the batch is ready to be committed.
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum BatchStatus {
43    /// The batch can accept more values.
44    Pending,
45    /// The batch is full and should be committed.
46    Ready,
47}
48
49/// Handlers implement the logic for a given indexing pipeline: How to process checkpoint data (by
50/// implementing [Processor]) into rows for their table, and how to write those rows to the database.
51///
52/// The handler is also responsible for tuning the various parameters of the pipeline (provided as
53/// associated values). Reasonable defaults have been chosen to balance concurrency with memory
54/// usage, but each handle may choose to override these defaults, e.g.
55///
56/// - Handlers that produce many small rows may wish to increase their batch/chunk/max-pending
57///   sizes).
58/// - Handlers that do more work during processing may wish to increase their fanout so more of it
59///   can be done concurrently, to preserve throughput.
60///
61/// Concurrent handlers can only be used in concurrent pipelines, where checkpoint data is
62/// processed and committed out-of-order and a watermark table is kept up-to-date with the latest
63/// checkpoint below which all data has been committed.
64///
65/// Back-pressure is handled through the `MAX_PENDING_SIZE` constant -- if more than this many rows
66/// build up, the collector will stop accepting new checkpoints, which will eventually propagate
67/// back to the ingestion service.
68#[async_trait]
69pub trait Handler: Processor {
70    type Store: ConcurrentStore;
71    type Batch: Default + Send + Sync + 'static;
72
73    /// If at least this many rows are pending, the committer will commit them eagerly.
74    const MIN_EAGER_ROWS: usize = 50;
75
76    /// If there are more than this many rows pending, the committer applies backpressure.
77    const MAX_PENDING_ROWS: usize = 5000;
78
79    /// The maximum number of watermarks that can show up in a single batch.
80    /// This limit exists to deal with pipelines that produce no data for a majority of
81    /// checkpoints -- the size of these pipeline's batches will be dominated by watermark updates.
82    const MAX_WATERMARK_UPDATES: usize = 10_000;
83
84    /// Add values from the iterator to the batch. The implementation may take all, some, or none
85    /// of the values from the iterator by calling `.next()`.
86    ///
87    /// Returns `BatchStatus::Ready` if the batch is full and should be committed,
88    /// or `BatchStatus::Pending` if the batch can accept more values.
89    ///
90    /// Note: The handler can signal batch readiness via `BatchStatus::Ready`, but the framework
91    /// may also decide to commit a batch based on the trait parameters above.
92    fn batch(
93        &self,
94        batch: &mut Self::Batch,
95        values: &mut std::vec::IntoIter<Self::Value>,
96    ) -> BatchStatus;
97
98    /// Commit the batch to the database, returning the number of rows affected.
99    async fn commit<'a>(
100        &self,
101        batch: &Self::Batch,
102        conn: &mut <Self::Store as Store>::Connection<'a>,
103    ) -> anyhow::Result<usize>;
104
105    /// Clean up data between checkpoints `_from` and `_to_exclusive` (exclusive) in the database, returning
106    /// the number of rows affected. This function is optional, and defaults to not pruning at all.
107    async fn prune<'a>(
108        &self,
109        _from: u64,
110        _to_exclusive: u64,
111        _conn: &mut <Self::Store as Store>::Connection<'a>,
112    ) -> anyhow::Result<usize> {
113        Ok(0)
114    }
115}
116
117/// Configuration for a concurrent pipeline
118#[derive(Serialize, Deserialize, Debug, Clone, Default)]
119pub struct ConcurrentConfig {
120    /// Configuration for the writer, that makes forward progress.
121    pub committer: CommitterConfig,
122
123    /// Per-pipeline ingestion overrides.
124    pub ingestion: IngestionConfig,
125
126    /// Configuration for the pruner, that deletes old data.
127    pub pruner: Option<PrunerConfig>,
128
129    /// Processor concurrency. Defaults to adaptive scaling up to the number of CPUs.
130    pub fanout: Option<ConcurrencyConfig>,
131
132    /// Override for `Handler::MIN_EAGER_ROWS` (eager batch threshold).
133    pub min_eager_rows: Option<usize>,
134
135    /// Override for `Handler::MAX_PENDING_ROWS` (backpressure threshold).
136    pub max_pending_rows: Option<usize>,
137
138    /// Override for `Handler::MAX_WATERMARK_UPDATES` (watermarks per batch cap).
139    pub max_watermark_updates: Option<usize>,
140
141    /// Size of the channel between the processor and collector.
142    pub processor_channel_size: Option<usize>,
143
144    /// Size of the channel between the collector and committer.
145    pub collector_channel_size: Option<usize>,
146
147    /// Size of the channel between the committer and the watermark updater.
148    pub committer_channel_size: Option<usize>,
149}
150
151#[derive(Serialize, Deserialize, Debug, Clone)]
152pub struct PrunerConfig {
153    /// How often the pruner should check whether there is any data to prune, in milliseconds.
154    pub interval_ms: u64,
155
156    /// How long to wait after the reader low watermark was set, until it is safe to prune up until
157    /// this new watermark, in milliseconds.
158    pub delay_ms: u64,
159
160    /// How much data to keep, this is measured in checkpoints.
161    pub retention: u64,
162
163    /// The maximum range to try and prune in one request, measured in checkpoints.
164    pub max_chunk_size: u64,
165
166    /// The max number of tasks to run in parallel for pruning.
167    pub prune_concurrency: u64,
168}
169
170/// Values ready to be written to the database. This is an internal type used to communicate
171/// between the collector and the committer parts of the pipeline.
172///
173/// Values inside each batch may or may not be from the same checkpoint. Values in the same
174/// checkpoint can also be split across multiple batches.
175struct BatchedRows<H: Handler> {
176    /// The batch to write
177    batch: H::Batch,
178    /// Number of rows in the batch
179    batch_len: usize,
180    /// Proportions of all the watermarks that are represented in this chunk
181    watermark: Vec<WatermarkPart>,
182}
183
184impl<H, V> BatchedRows<H>
185where
186    H: Handler<Batch = Vec<V>, Value = V>,
187{
188    #[cfg(test)]
189    pub fn from_vec(batch: Vec<V>, watermark: Vec<WatermarkPart>) -> Self {
190        let batch_len = batch.len();
191        Self {
192            batch,
193            batch_len,
194            watermark,
195        }
196    }
197}
198
199impl PrunerConfig {
200    pub fn interval(&self) -> Duration {
201        Duration::from_millis(self.interval_ms)
202    }
203
204    pub fn delay(&self) -> Duration {
205        Duration::from_millis(self.delay_ms)
206    }
207}
208
209impl Default for PrunerConfig {
210    fn default() -> Self {
211        Self {
212            interval_ms: 300_000,
213            delay_ms: 120_000,
214            retention: 4_000_000,
215            max_chunk_size: 2_000,
216            prune_concurrency: 1,
217        }
218    }
219}
220
221/// Start a new concurrent (out-of-order) indexing pipeline served by the handler, `H`. Starting
222/// strictly after the `watermark` (or from the beginning if no watermark was provided).
223///
224/// Each pipeline consists of a processor task which takes checkpoint data and breaks it down into
225/// rows, ready for insertion, a collector which batches those rows into an appropriate size for
226/// the database, a committer which writes the rows out concurrently, and a watermark task to
227/// update the high watermark.
228///
229/// Committing is performed out-of-order: the pipeline may write out checkpoints out-of-order,
230/// either because it received the checkpoints out-of-order or because of variance in processing
231/// time.
232///
233/// The pipeline also maintains a row in the `watermarks` table for the pipeline which tracks the
234/// watermark below which all data has been committed (modulo pruning).
235///
236/// Checkpoint data is fed into the pipeline through the `checkpoint_rx` channel, and internal
237/// channels are created to communicate between its various components. The pipeline will shutdown
238/// if any of its input or output channels close, any of its independent tasks fail, or if it is
239/// signalled to shutdown through the returned service handle.
240pub(crate) fn pipeline<H: Handler>(
241    handler: H,
242    next_checkpoint: u64,
243    config: ConcurrentConfig,
244    store: H::Store,
245    task: Option<Task>,
246    checkpoint_rx: mpsc::Receiver<Arc<CheckpointEnvelope>>,
247    metrics: Arc<IndexerMetrics>,
248) -> Service {
249    info!(
250        pipeline = H::NAME,
251        "Starting pipeline with config: {config:#?}",
252    );
253
254    let ConcurrentConfig {
255        committer: committer_config,
256        ingestion: _,
257        pruner: pruner_config,
258        fanout,
259        min_eager_rows,
260        max_pending_rows,
261        max_watermark_updates,
262        processor_channel_size,
263        collector_channel_size,
264        committer_channel_size,
265    } = config;
266
267    let concurrency = fanout.unwrap_or(ConcurrencyConfig::Adaptive {
268        initial: 1,
269        min: 1,
270        max: num_cpus::get().max(1),
271        dead_band: None,
272    });
273    let min_eager_rows = min_eager_rows.unwrap_or(H::MIN_EAGER_ROWS);
274    let max_pending_rows = max_pending_rows.unwrap_or(H::MAX_PENDING_ROWS);
275    let max_watermark_updates = max_watermark_updates.unwrap_or(H::MAX_WATERMARK_UPDATES);
276
277    let processor_channel_size = processor_channel_size
278        .unwrap_or_else(|| num_cpus::get() / 2)
279        .max(1);
280    let (processor_tx, collector_rx) = mpsc::channel(processor_channel_size);
281
282    let collector_channel_size = collector_channel_size
283        .unwrap_or_else(|| num_cpus::get() / 2)
284        .max(1);
285    //docs::#buff (see docs/content/guides/developer/advanced/custom-indexer.mdx)
286    let (collector_tx, committer_rx) = mpsc::channel(collector_channel_size);
287    //docs::/#buff
288    let committer_channel_size = committer_channel_size.unwrap_or_else(num_cpus::get).max(1);
289    let (committer_tx, watermark_rx) = mpsc::channel(committer_channel_size);
290    let main_reader_lo = Arc::new(SetOnce::new());
291
292    let handler = Arc::new(handler);
293
294    let s_processor = processor(
295        handler.clone(),
296        checkpoint_rx,
297        processor_tx,
298        metrics.clone(),
299        concurrency,
300        store.clone(),
301    );
302
303    let s_collector = collector::<H>(
304        handler.clone(),
305        committer_config.clone(),
306        collector_rx,
307        collector_tx,
308        main_reader_lo.clone(),
309        metrics.clone(),
310        min_eager_rows,
311        max_pending_rows,
312        max_watermark_updates,
313    );
314
315    let s_committer = committer::<H>(
316        handler.clone(),
317        committer_config.clone(),
318        committer_rx,
319        committer_tx,
320        store.clone(),
321        metrics.clone(),
322    );
323
324    let s_commit_watermark = commit_watermark::<H>(
325        next_checkpoint,
326        committer_config,
327        watermark_rx,
328        store.clone(),
329        task.as_ref().map(|t| t.task.clone()),
330        metrics.clone(),
331    );
332
333    let s_track_reader_lo = track_main_reader_lo::<H>(
334        main_reader_lo.clone(),
335        task.as_ref().map(|t| t.reader_interval),
336        store.clone(),
337    );
338
339    let s_reader_watermark =
340        reader_watermark::<H>(pruner_config.clone(), store.clone(), metrics.clone());
341
342    let s_pruner = pruner(handler, pruner_config, store, metrics);
343
344    s_processor
345        .merge(s_collector)
346        .merge(s_committer)
347        .merge(s_commit_watermark)
348        .attach(s_track_reader_lo)
349        .attach(s_reader_watermark)
350        .attach(s_pruner)
351}
352
353#[cfg(test)]
354mod tests {
355    use std::sync::Arc;
356    use std::time::Duration;
357
358    use prometheus::Registry;
359    use sui_types::digests::CheckpointDigest;
360    use tokio::sync::mpsc;
361    use tokio::time::timeout;
362
363    use crate::FieldCount;
364    use crate::metrics::IndexerMetrics;
365    use crate::mocks::store::FallibleMockConnection;
366    use crate::mocks::store::FallibleMockStore;
367    use crate::pipeline::Processor;
368    use crate::types::full_checkpoint_content::Checkpoint;
369    use crate::types::test_checkpoint_data_builder::TestCheckpointBuilder;
370
371    use super::*;
372
373    const TEST_TIMEOUT: Duration = Duration::from_secs(60);
374    const TEST_SUBSCRIBER_CHANNEL_SIZE: usize = 3; // Critical for back-pressure testing calculations
375
376    #[derive(Clone, Debug, FieldCount)]
377    struct TestValue {
378        checkpoint: u64,
379        data: u64,
380    }
381
382    struct DataPipeline;
383
384    #[async_trait]
385    impl Processor for DataPipeline {
386        const NAME: &'static str = "test_handler";
387        type Value = TestValue;
388
389        async fn process(&self, checkpoint: &Arc<Checkpoint>) -> anyhow::Result<Vec<Self::Value>> {
390            let cp_num = checkpoint.summary.sequence_number;
391
392            // Every checkpoint will come with 2 processed values
393            Ok(vec![
394                TestValue {
395                    checkpoint: cp_num,
396                    data: cp_num * 10 + 1,
397                },
398                TestValue {
399                    checkpoint: cp_num,
400                    data: cp_num * 10 + 2,
401                },
402            ])
403        }
404    }
405
406    #[async_trait]
407    impl Handler for DataPipeline {
408        type Store = FallibleMockStore;
409        type Batch = Vec<TestValue>;
410
411        const MIN_EAGER_ROWS: usize = 1000; // High value to disable eager batching
412        const MAX_PENDING_ROWS: usize = 4; // Small value to trigger back pressure quickly
413        const MAX_WATERMARK_UPDATES: usize = 1; // Each batch will have 1 checkpoint for an ease of testing.
414
415        fn batch(
416            &self,
417            batch: &mut Self::Batch,
418            values: &mut std::vec::IntoIter<Self::Value>,
419        ) -> BatchStatus {
420            // Take all values
421            batch.extend(values);
422            BatchStatus::Pending
423        }
424
425        async fn commit<'a>(
426            &self,
427            batch: &Self::Batch,
428            conn: &mut FallibleMockConnection<'a>,
429        ) -> anyhow::Result<usize> {
430            // Group values by checkpoint
431            let mut grouped: std::collections::HashMap<u64, Vec<u64>> =
432                std::collections::HashMap::new();
433            for value in batch {
434                grouped
435                    .entry(value.checkpoint)
436                    .or_default()
437                    .push(value.data);
438            }
439
440            // Commit all data at once
441            conn.0.commit_bulk_data(DataPipeline::NAME, grouped).await
442        }
443
444        async fn prune<'a>(
445            &self,
446            from: u64,
447            to_exclusive: u64,
448            conn: &mut FallibleMockConnection<'a>,
449        ) -> anyhow::Result<usize> {
450            conn.0.prune_data(DataPipeline::NAME, from, to_exclusive)
451        }
452    }
453
454    struct TestSetup {
455        store: FallibleMockStore,
456        checkpoint_tx: mpsc::Sender<Arc<CheckpointEnvelope>>,
457        #[allow(unused)]
458        pipeline: Service,
459    }
460
461    impl TestSetup {
462        async fn new(
463            config: ConcurrentConfig,
464            store: FallibleMockStore,
465            next_checkpoint: u64,
466        ) -> Self {
467            let (checkpoint_tx, checkpoint_rx) = mpsc::channel(TEST_SUBSCRIBER_CHANNEL_SIZE);
468            let metrics = IndexerMetrics::new(None, &Registry::default());
469
470            let pipeline = pipeline(
471                DataPipeline,
472                next_checkpoint,
473                config,
474                store.clone(),
475                None,
476                checkpoint_rx,
477                metrics,
478            );
479
480            Self {
481                store,
482                checkpoint_tx,
483                pipeline,
484            }
485        }
486
487        async fn send_checkpoint(&self, checkpoint: u64) -> anyhow::Result<()> {
488            let checkpoint_envelope = Arc::new(CheckpointEnvelope {
489                checkpoint: Arc::new(
490                    TestCheckpointBuilder::new(checkpoint)
491                        .with_epoch(1)
492                        .with_network_total_transactions(checkpoint * 2)
493                        .with_timestamp_ms(1000000000 + checkpoint * 1000)
494                        .build_checkpoint(),
495                ),
496                chain_id: CheckpointDigest::new([1; 32]).into(),
497            });
498            self.checkpoint_tx.send(checkpoint_envelope).await?;
499            Ok(())
500        }
501
502        async fn send_checkpoint_with_timeout(
503            &self,
504            checkpoint: u64,
505            timeout_duration: Duration,
506        ) -> anyhow::Result<()> {
507            timeout(timeout_duration, self.send_checkpoint(checkpoint)).await?
508        }
509
510        async fn send_checkpoint_expect_timeout(
511            &self,
512            checkpoint: u64,
513            timeout_duration: Duration,
514        ) {
515            timeout(timeout_duration, self.send_checkpoint(checkpoint))
516                .await
517                .unwrap_err(); // Panics if send succeeds
518        }
519    }
520
521    #[tokio::test]
522    async fn test_e2e_pipeline() {
523        let config = ConcurrentConfig {
524            pruner: Some(PrunerConfig {
525                interval_ms: 5_000, // Long interval to test states before pruning
526                delay_ms: 100,      // Short delay for faster tests
527                retention: 3,       // Keep only 3 checkpoints
528                ..Default::default()
529            }),
530            ..Default::default()
531        };
532        let store = FallibleMockStore::default();
533        let setup = TestSetup::new(config, store, 0).await;
534
535        // Send initial checkpoints
536        for i in 0..3 {
537            setup
538                .send_checkpoint_with_timeout(i, Duration::from_millis(200))
539                .await
540                .unwrap();
541        }
542
543        // Verify all initial data is available (before any pruning)
544        for i in 0..3 {
545            let data = setup
546                .store
547                .wait_for_data(DataPipeline::NAME, i, TEST_TIMEOUT)
548                .await;
549            assert_eq!(data, vec![i * 10 + 1, i * 10 + 2]);
550        }
551
552        // Add more checkpoints to trigger pruning
553        for i in 3..6 {
554            setup
555                .send_checkpoint_with_timeout(i, Duration::from_millis(200))
556                .await
557                .unwrap();
558        }
559
560        // Verify data is still available BEFORE pruning kicks in
561        // With long pruning interval (5s), we can safely verify data without race conditions
562        for i in 0..6 {
563            let data = setup
564                .store
565                .wait_for_data(DataPipeline::NAME, i, Duration::from_secs(1))
566                .await;
567            assert_eq!(data, vec![i * 10 + 1, i * 10 + 2]);
568        }
569
570        // Wait for pruning to occur. The pruner and reader_watermark tasks both run on
571        // the same interval, so poll until the pruner has caught up rather instead of using a
572        // fixed sleep.
573        let pruning_deadline = Duration::from_secs(15);
574        let start = tokio::time::Instant::now();
575        loop {
576            let pruned = {
577                let data = setup.store.data.get(DataPipeline::NAME).unwrap();
578                !data.contains_key(&0) && !data.contains_key(&1) && !data.contains_key(&2)
579            };
580            if pruned {
581                break;
582            }
583            assert!(
584                start.elapsed() < pruning_deadline,
585                "Timed out waiting for pruning to occur"
586            );
587            tokio::time::sleep(Duration::from_millis(100)).await;
588        }
589
590        // Verify recent checkpoints are still available
591        {
592            let data = setup.store.data.get(DataPipeline::NAME).unwrap();
593            assert!(data.contains_key(&3));
594            assert!(data.contains_key(&4));
595            assert!(data.contains_key(&5));
596        };
597    }
598
599    #[tokio::test]
600    async fn test_e2e_pipeline_without_pruning() {
601        let config = ConcurrentConfig {
602            pruner: None,
603            ..Default::default()
604        };
605        let store = FallibleMockStore::default();
606        let setup = TestSetup::new(config, store, 0).await;
607
608        // Send several checkpoints
609        for i in 0..10 {
610            setup
611                .send_checkpoint_with_timeout(i, Duration::from_millis(200))
612                .await
613                .unwrap();
614        }
615
616        // Wait for all data to be processed and committed
617        let watermark = setup
618            .store
619            .wait_for_watermark(DataPipeline::NAME, 9, TEST_TIMEOUT)
620            .await;
621
622        // Verify ALL data was processed correctly (no pruning)
623        for i in 0..10 {
624            let data = setup
625                .store
626                .wait_for_data(DataPipeline::NAME, i, Duration::from_secs(1))
627                .await;
628            assert_eq!(data, vec![i * 10 + 1, i * 10 + 2]);
629        }
630
631        // Verify watermark progression
632        assert_eq!(watermark.checkpoint_hi_inclusive, Some(9));
633        assert_eq!(watermark.tx_hi, 18); // 9 * 2
634        assert_eq!(watermark.timestamp_ms_hi_inclusive, 1000009000); // 1000000000 + 9 * 1000
635
636        // Verify no data was pruned - all 10 checkpoints should still exist
637        let total_checkpoints = {
638            let data = setup.store.data.get(DataPipeline::NAME).unwrap();
639            data.len()
640        };
641        assert_eq!(total_checkpoints, 10);
642    }
643
644    #[tokio::test]
645    async fn test_out_of_order_processing() {
646        let config = ConcurrentConfig::default();
647        let store = FallibleMockStore::default();
648        let setup = TestSetup::new(config, store, 0).await;
649
650        // Send checkpoints out of order
651        let checkpoints = vec![2, 0, 4, 1, 3];
652        for cp in checkpoints {
653            setup
654                .send_checkpoint_with_timeout(cp, Duration::from_millis(200))
655                .await
656                .unwrap();
657        }
658
659        // Wait for all data to be processed
660        setup
661            .store
662            .wait_for_watermark(DataPipeline::NAME, 4, Duration::from_secs(5))
663            .await;
664
665        // Verify all checkpoints were processed correctly despite out-of-order arrival
666        for i in 0..5 {
667            let data = setup
668                .store
669                .wait_for_data(DataPipeline::NAME, i, Duration::from_secs(1))
670                .await;
671            assert_eq!(data, vec![i * 10 + 1, i * 10 + 2]);
672        }
673    }
674
675    #[tokio::test]
676    async fn test_watermark_progression_with_gaps() {
677        let config = ConcurrentConfig::default();
678        let store = FallibleMockStore::default();
679        let setup = TestSetup::new(config, store, 0).await;
680
681        // Send checkpoints with a gap (0, 1, 3, 4) - missing checkpoint 2
682        for cp in [0, 1, 3, 4] {
683            setup
684                .send_checkpoint_with_timeout(cp, Duration::from_millis(200))
685                .await
686                .unwrap();
687        }
688
689        // Wait for processing
690        tokio::time::sleep(Duration::from_secs(1)).await;
691
692        // Watermark should only progress to 1 (can't progress past the gap)
693        let watermark = setup.store.watermark(DataPipeline::NAME).unwrap();
694        assert_eq!(watermark.checkpoint_hi_inclusive, Some(1));
695
696        // Now send the missing checkpoint 2
697        setup
698            .send_checkpoint_with_timeout(2, Duration::from_millis(200))
699            .await
700            .unwrap();
701
702        // Now watermark should progress to 4
703        let watermark = setup
704            .store
705            .wait_for_watermark(DataPipeline::NAME, 4, TEST_TIMEOUT)
706            .await;
707        assert_eq!(watermark.checkpoint_hi_inclusive, Some(4));
708    }
709
710    // ==================== BACK-PRESSURE TESTING ====================
711
712    #[tokio::test]
713    async fn test_back_pressure_collector_max_pending_rows() {
714        // Pipeline Diagram - Collector Back Pressure via MAX_PENDING_ROWS:
715        //
716        // ┌────────────┐    ┌────────────┐    ┌────────────┐    ┌────────────┐
717        // │ Checkpoint │ ─► │ Processor  │ ─► │ Collector  │ ─► │ Committer  │
718        // │   Input    │    │ (fanout=2) │    │            │    │            │
719        // └────────────┘    └────────────┘    └[BOTTLENECK]┘    └────────────┘
720        //                │                 │                 │
721        //              [●●●]           [●●●●●●●]         [●●●●●●]
722        //            buffer: 3         buffer: 7         buffer: 6
723        //
724        // BOTTLENECK: Collector stops accepting when pending rows ≥ MAX_PENDING_ROWS (4)
725
726        let config = ConcurrentConfig {
727            committer: CommitterConfig {
728                collect_interval_ms: 5_000, // Long interval to prevent timer-driven collection
729                write_concurrency: 1,
730                ..Default::default()
731            },
732            fanout: Some(ConcurrencyConfig::Fixed { value: 2 }),
733            processor_channel_size: Some(7),
734            collector_channel_size: Some(6),
735            ..Default::default()
736        };
737        let store = FallibleMockStore::default();
738        let setup = TestSetup::new(config, store, 0).await;
739
740        // Wait for initial setup
741        tokio::time::sleep(Duration::from_millis(200)).await;
742
743        // Pipeline capacity analysis with collector back pressure:
744        // Configuration: MAX_PENDING_ROWS=4, fanout=2
745        //
746        // Channel and task breakdown:
747        // - Checkpoint->Processor channel: 3 slots (TEST_SUBSCRIBER_CHANNEL_SIZE)
748        // - Processor tasks: 2 tasks (fanout=2)
749        // - Processor->Collector channel: 7 slots (processor_channel_size=7)
750        // - Collector pending: 2 checkpoints × 2 values = 4 values (hits MAX_PENDING_ROWS=4)
751        //
752        // Total capacity: 3 + 2 + 7 + 2 = 14 checkpoints
753
754        // Fill pipeline to capacity - these should all succeed
755        for i in 0..14 {
756            setup
757                .send_checkpoint_with_timeout(i, Duration::from_millis(200))
758                .await
759                .unwrap();
760        }
761
762        // Checkpoint 14 should block due to MAX_PENDING_ROWS back pressure
763        setup
764            .send_checkpoint_expect_timeout(14, Duration::from_millis(200))
765            .await;
766
767        // Allow pipeline to drain by sending the blocked checkpoint with longer timeout
768        setup
769            .send_checkpoint_with_timeout(14, TEST_TIMEOUT)
770            .await
771            .unwrap();
772
773        // Verify data was processed correctly
774        let data = setup
775            .store
776            .wait_for_data(DataPipeline::NAME, 0, TEST_TIMEOUT)
777            .await;
778        assert_eq!(data, vec![1, 2]);
779    }
780
781    #[tokio::test]
782    async fn test_back_pressure_committer_slow_commits() {
783        // Pipeline Diagram - Committer Back Pressure via Slow Database Commits:
784        //
785        // ┌────────────┐    ┌────────────┐    ┌────────────┐    ┌────────────┐
786        // │ Checkpoint │ ─► │ Processor  │ ─► │ Collector  │ ─► │ Committer  │
787        // │   Input    │    │ (fanout=2) │    │            │    │🐌 10s Delay│
788        // └────────────┘    └────────────┘    └────────────┘    └[BOTTLENECK]┘
789        //                │                 │                 │
790        //              [●●●]           [●●●●●●●]          [●●●●●●]
791        //            buffer: 3    proc_chan: 7       coll_chan: 6
792        //
793        // BOTTLENECK: Committer with 10s delay blocks entire pipeline
794
795        let config = ConcurrentConfig {
796            committer: CommitterConfig {
797                write_concurrency: 1, // Single committer for deterministic blocking
798                // MIN_EAGER_ROWS is 1000 and MAX_PENDING_ROWS is 4, so
799                // this test relies on the collect interval tip to force
800                // batch flushing.
801                collect_interval_ms: 10,
802                ..Default::default()
803            },
804            fanout: Some(ConcurrencyConfig::Fixed { value: 2 }),
805            processor_channel_size: Some(7),
806            collector_channel_size: Some(6),
807            ..Default::default()
808        };
809        let store = FallibleMockStore::default().with_commit_delay(10_000); // 10 seconds delay
810        let setup = TestSetup::new(config, store, 0).await;
811
812        // Pipeline capacity analysis with slow commits:
813        // Configuration: fanout=2, write_concurrency=1
814        //
815        // Channel and task breakdown:
816        // - Checkpoint->Processor channel: 3 slots (TEST_SUBSCRIBER_CHANNEL_SIZE)
817        // - Processor tasks: 2 tasks (fanout=2)
818        // - Processor->Collector channel: 7 slots (processor_channel_size=7)
819        // - Collector->Committer channel: 6 slots (collector_channel_size=6)
820        // - Committer task: 1 task (blocked by slow commit)
821        //
822        // Total theoretical capacity: 3 + 2 + 7 + 6 + 1 = 19 checkpoints
823
824        // Fill pipeline to theoretical capacity - these should all succeed
825        for i in 0..19 {
826            setup
827                .send_checkpoint_with_timeout(i, Duration::from_millis(100))
828                .await
829                .unwrap();
830        }
831
832        // Find the actual back pressure point
833        // Due to non-determinism in collector's tokio::select!, the collector may consume
834        // up to 2 checkpoints (filling MAX_PENDING_ROWS=4) before applying back pressure.
835        // This means back pressure occurs somewhere in range 19-21.
836        let mut back_pressure_checkpoint = None;
837        for checkpoint in 19..22 {
838            if setup
839                .send_checkpoint_with_timeout(checkpoint, Duration::from_millis(100))
840                .await
841                .is_err()
842            {
843                back_pressure_checkpoint = Some(checkpoint);
844                break;
845            }
846        }
847        assert!(
848            back_pressure_checkpoint.is_some(),
849            "Back pressure should occur between checkpoints 19-21"
850        );
851
852        // Verify that some data has been processed (pipeline is working)
853        setup
854            .store
855            .wait_for_any_data(DataPipeline::NAME, TEST_TIMEOUT)
856            .await;
857
858        // Allow pipeline to drain by sending the blocked checkpoint with longer timeout
859        setup
860            .send_checkpoint_with_timeout(back_pressure_checkpoint.unwrap(), TEST_TIMEOUT)
861            .await
862            .unwrap();
863    }
864
865    // ==================== FAILURE TESTING ====================
866
867    #[tokio::test]
868    async fn test_commit_failure_retry() {
869        let config = ConcurrentConfig::default();
870        let store = FallibleMockStore::default().with_commit_failures(2); // Fail 2 times, then succeed
871        let setup = TestSetup::new(config, store, 0).await;
872
873        // Send a checkpoint
874        setup
875            .send_checkpoint_with_timeout(0, Duration::from_millis(200))
876            .await
877            .unwrap();
878
879        // Should eventually succeed despite initial commit failures
880        setup
881            .store
882            .wait_for_watermark(DataPipeline::NAME, 0, TEST_TIMEOUT)
883            .await;
884
885        // Verify data was eventually committed
886        let data = setup
887            .store
888            .wait_for_data(DataPipeline::NAME, 0, Duration::from_secs(1))
889            .await;
890        assert_eq!(data, vec![1, 2]);
891    }
892
893    #[tokio::test]
894    async fn test_prune_failure_retry() {
895        let config = ConcurrentConfig {
896            pruner: Some(PrunerConfig {
897                interval_ms: 2000, // 2 seconds interval for testing
898                delay_ms: 100,     // Short delay
899                retention: 2,      // Keep only 2 checkpoints
900                ..Default::default()
901            }),
902            ..Default::default()
903        };
904
905        // Configure prune failures for range [0, 2) - fail twice then succeed
906        let store = FallibleMockStore::default().with_prune_failures(0, 2, 1);
907        let setup = TestSetup::new(config, store, 0).await;
908
909        // Send enough checkpoints to trigger pruning
910        for i in 0..4 {
911            setup
912                .send_checkpoint_with_timeout(i, Duration::from_millis(200))
913                .await
914                .unwrap();
915        }
916
917        // Verify data is still available BEFORE pruning kicks in
918        // With long pruning interval (5s), we can safely verify data without race conditions
919        for i in 0..4 {
920            let data = setup
921                .store
922                .wait_for_data(DataPipeline::NAME, i, Duration::from_secs(1))
923                .await;
924            assert_eq!(data, vec![i * 10 + 1, i * 10 + 2]);
925        }
926
927        // Wait for first pruning attempt (should fail) and verify no data has been pruned
928        setup
929            .store
930            .wait_for_prune_attempts(0, 2, 1, TEST_TIMEOUT)
931            .await;
932        {
933            let data = setup.store.data.get(DataPipeline::NAME).unwrap();
934            for i in 0..4 {
935                assert!(data.contains_key(&i));
936            }
937        };
938
939        // Wait for second pruning attempt (should succeed)
940        setup
941            .store
942            .wait_for_prune_attempts(0, 2, 2, TEST_TIMEOUT)
943            .await;
944        {
945            let data = setup.store.data.get(DataPipeline::NAME).unwrap();
946            // Verify recent checkpoints are still available
947            assert!(data.contains_key(&2));
948            assert!(data.contains_key(&3));
949
950            // Verify old checkpoints are pruned
951            assert!(!data.contains_key(&0));
952            assert!(!data.contains_key(&1));
953        };
954    }
955
956    #[tokio::test]
957    async fn test_reader_watermark_failure_retry() {
958        let config = ConcurrentConfig {
959            pruner: Some(PrunerConfig {
960                interval_ms: 100, // Fast interval for testing
961                delay_ms: 100,    // Short delay
962                retention: 3,     // Keep 3 checkpoints
963                ..Default::default()
964            }),
965            ..Default::default()
966        };
967
968        // Configure reader watermark failures - fail 2 times then succeed
969        let store = FallibleMockStore::default().with_reader_watermark_failures(2);
970        let setup = TestSetup::new(config, store, 0).await;
971
972        // Send checkpoints to trigger reader watermark updates
973        for i in 0..6 {
974            setup
975                .send_checkpoint_with_timeout(i, Duration::from_millis(200))
976                .await
977                .unwrap();
978        }
979
980        // Wait for processing to complete
981        setup
982            .store
983            .wait_for_watermark(DataPipeline::NAME, 5, TEST_TIMEOUT)
984            .await;
985
986        // Wait for reader watermark task to attempt updates (with failures and retries)
987        tokio::time::sleep(Duration::from_secs(2)).await;
988
989        // Verify that reader watermark was eventually updated despite failures
990        let watermark = setup.store.watermark(DataPipeline::NAME).unwrap();
991        assert_eq!(watermark.reader_lo, 3);
992    }
993
994    #[tokio::test]
995    async fn test_database_connection_failure_retry() {
996        let config = ConcurrentConfig::default();
997        let store = FallibleMockStore::default().with_connection_failures(2); // Fail 2 times, then succeed
998        let setup = TestSetup::new(config, store, 0).await;
999
1000        // Send a checkpoint
1001        setup
1002            .send_checkpoint_with_timeout(0, Duration::from_millis(200))
1003            .await
1004            .unwrap();
1005
1006        // Should eventually succeed despite initial failures
1007        setup
1008            .store
1009            .wait_for_watermark(DataPipeline::NAME, 0, TEST_TIMEOUT)
1010            .await;
1011
1012        // Verify data was eventually committed
1013        let data = setup
1014            .store
1015            .wait_for_data(DataPipeline::NAME, 0, TEST_TIMEOUT)
1016            .await;
1017        assert_eq!(data, vec![1, 2]);
1018    }
1019}