sui_indexer_alt_framework/ingestion/
mod.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::num::NonZeroUsize;
5use std::sync::Arc;
6use std::time::Duration;
7
8use prometheus::Registry;
9use serde::Deserialize;
10use serde::Serialize;
11use sui_futures::service::Service;
12use tokio::sync::mpsc;
13use tracing::warn;
14
15use crate::cohort::DEFAULT_MIN_COHORT_BOUNDARY;
16pub use crate::config::ConcurrencyConfig as IngestConcurrencyConfig;
17use crate::ingestion::broadcaster::broadcaster;
18use crate::ingestion::error::Error;
19use crate::ingestion::error::Result;
20use crate::ingestion::ingestion_client::CheckpointEnvelope;
21use crate::ingestion::ingestion_client::IngestionClient;
22use crate::ingestion::ingestion_client::IngestionClientArgs;
23use crate::ingestion::ingestion_client::retry_transient_with_slow_monitor;
24use crate::ingestion::streaming_client::CheckpointStreamingClient;
25use crate::ingestion::streaming_client::GrpcStreamingClient;
26use crate::ingestion::streaming_client::StreamingClientArgs;
27use crate::metrics::IngestionMetrics;
28
29/// Type alias for a shared [`CheckpointStreamingClient`] trait object,
30/// the form [`IngestionService`] stores and the broadcaster consumes.
31/// `Arc`'d (rather than `Box`'d) because `CheckpointStreamingClient`'s
32/// methods take `&self`, so the client can be cheaply cloned and shared
33/// across owners -- each cohort's [`IngestionService`] gets its own clone.
34/// The `Send + Sync` bounds let it move across task boundaries and be
35/// shared behind a reference when an enclosing [`IngestionService`] is
36/// held across threads.
37pub type ArcStreamingClient = Arc<dyn CheckpointStreamingClient + Send + Sync>;
38
39mod broadcaster;
40mod byte_count;
41pub(crate) mod decode;
42pub mod error;
43pub mod ingestion_client;
44mod rpc_client;
45pub mod store_client;
46pub mod streaming_client;
47#[cfg(test)]
48mod test_utils;
49
50pub(crate) const MAX_GRPC_MESSAGE_SIZE_BYTES: usize = 128 * 1024 * 1024;
51
52/// Combined arguments for both ingestion and streaming clients.
53/// This is a convenience wrapper that flattens both argument types.
54#[derive(clap::Args, Clone, Debug, Default)]
55pub struct ClientArgs {
56    #[clap(flatten)]
57    pub ingestion: IngestionClientArgs,
58
59    #[clap(flatten)]
60    pub streaming: StreamingClientArgs,
61}
62
63#[derive(Serialize, Deserialize, Debug, Clone)]
64#[serde(default)]
65pub struct IngestionConfig {
66    /// Concurrency control for checkpoint ingestion. A plain integer gives fixed concurrency;
67    /// an object with `initial`, `min`, and `max` fields enables adaptive concurrency that adjusts
68    /// based on subscriber channel fill fraction.
69    pub ingest_concurrency: IngestConcurrencyConfig,
70
71    /// Polling interval to retry fetching checkpoints that do not exist, in milliseconds.
72    pub retry_interval_ms: u64,
73
74    /// Initial number of checkpoints to process using ingestion after a streaming connection failure.
75    pub streaming_backoff_initial_batch_size: NonZeroUsize,
76
77    /// Maximum number of checkpoints to process using ingestion after repeated streaming connection failures.
78    pub streaming_backoff_max_batch_size: usize,
79
80    /// Timeout for streaming connection in milliseconds.
81    pub streaming_connection_timeout_ms: u64,
82
83    /// Timeout for streaming statement (peek/next) operations in milliseconds.
84    pub streaming_statement_timeout_ms: u64,
85
86    /// Minimum boundary (in checkpoints) for an ingestion cohort: the closest pipeline in a
87    /// cohort absorbs peers within `max(2 * its distance from the tip, this)`, so nearly
88    /// caught-up pipelines are not fragmented into many tiny cohorts. Defaults to 25,000
89    /// checkpoints when unset.
90    pub min_cohort_boundary: u64,
91}
92
93pub struct IngestionService {
94    config: IngestionConfig,
95    ingestion_client: IngestionClient,
96    streaming_client: Option<ArcStreamingClient>,
97    gap_fill_client: Option<IngestionClient>,
98    subscribers: Vec<mpsc::Sender<Arc<CheckpointEnvelope>>>,
99}
100
101/// Creates the [IngestionService]s that an indexer runs -- one per cohort of pipelines with
102/// similar distances from the network tip. The factory holds the clients services are made from;
103/// services are only constructed when requested through [Self::create], and every service is
104/// created the same way: it shares the factory's ingestion client (and so reports to the same
105/// metrics -- counters aggregate across services, while latest-checkpoint style gauges are labeled
106/// by cohort so concurrently-running services don't overwrite each other), and it gets its own
107/// clone of the streaming client, if there is one.
108pub(crate) struct IngestionFactory {
109    /// The ingestion configuration shared by every service this factory creates.
110    config: IngestionConfig,
111
112    /// The client shared by [Self::latest_checkpoint_number] and the services this factory
113    /// creates.
114    ingestion_client: IngestionClient,
115
116    /// Cloned into every service created, and used by [Self::latest_checkpoint_number].
117    streaming_client: Option<ArcStreamingClient>,
118
119    /// Rebound to each service's cohort, like the ingestion client.
120    gap_fill_client: Option<IngestionClient>,
121}
122
123impl IngestionConfig {
124    pub fn retry_interval(&self) -> Duration {
125        Duration::from_millis(self.retry_interval_ms)
126    }
127
128    pub fn streaming_connection_timeout(&self) -> Duration {
129        Duration::from_millis(self.streaming_connection_timeout_ms)
130    }
131
132    pub fn streaming_statement_timeout(&self) -> Duration {
133        Duration::from_millis(self.streaming_statement_timeout_ms)
134    }
135}
136
137impl IngestionService {
138    /// Create a new instance of the ingestion service, responsible for fetching checkpoints and
139    /// disseminating them to subscribers.
140    ///
141    /// - `args` specifies where to fetch checkpoints from.
142    /// - `config` specifies the various sizes and time limits for ingestion.
143    /// - `metrics_prefix` and `registry` are used to set up metrics for the service.
144    ///
145    /// After initialization, subscribers can be added using [Self::subscribe_bounded], and the
146    /// service is started with [Self::run], given a range of checkpoints to fetch (potentially
147    /// unbounded).
148    pub async fn new(
149        args: ClientArgs,
150        config: IngestionConfig,
151        metrics_prefix: Option<&str>,
152        registry: &Registry,
153    ) -> Result<Self> {
154        let metrics = IngestionMetrics::new(metrics_prefix, registry);
155        let (ingestion_client, streaming_client, gap_fill_client) =
156            clients_from_args(args, &config, metrics).await?;
157        Ok(Self::with_clients(
158            ingestion_client,
159            streaming_client,
160            gap_fill_client,
161            config,
162        ))
163    }
164
165    /// Construct an [`IngestionService`] from pre-built clients, bypassing [`ClientArgs`]-driven
166    /// construction.
167    ///
168    /// Callers that supply their own [`IngestionClientTrait`] / [`CheckpointStreamingClient`]
169    /// implementations — for example, when embedding the indexer in a fullnode that already has
170    /// checkpoint data on hand — use this constructor instead of [`Self::new`].
171    ///
172    /// If `gap_fill_client` is provided, it fetches the checkpoints between the service's
173    /// position and the start of a newly connected stream, instead of `ingestion_client`.
174    ///
175    /// [`IngestionClientTrait`]: crate::ingestion::ingestion_client::IngestionClientTrait
176    pub fn with_clients(
177        ingestion_client: IngestionClient,
178        streaming_client: Option<ArcStreamingClient>,
179        gap_fill_client: Option<IngestionClient>,
180        config: IngestionConfig,
181    ) -> Self {
182        Self {
183            config,
184            ingestion_client,
185            streaming_client,
186            gap_fill_client,
187            subscribers: Vec::new(),
188        }
189    }
190
191    /// The ingestion configuration this service was built with.
192    pub fn config(&self) -> &IngestionConfig {
193        &self.config
194    }
195
196    /// Add a new subscription backed by a bounded `mpsc` channel of the given capacity. The
197    /// channel itself is the backpressure signal: when this consumer falls behind, the channel
198    /// fills and the adaptive ingestion controller cuts fetch concurrency. Send blocks when the
199    /// channel is full.
200    ///
201    /// Callers typically pass `pipeline::IngestionConfig::subscriber_channel_size()`.
202    pub fn subscribe_bounded(&mut self, size: usize) -> mpsc::Receiver<Arc<CheckpointEnvelope>> {
203        let (tx, rx) = mpsc::channel(size);
204        self.subscribers.push(tx);
205        rx
206    }
207
208    /// Start the ingestion service as a background task, consuming it in the process.
209    ///
210    /// Checkpoints are fetched concurrently from the `checkpoints` iterator and pushed to
211    /// subscribers' channels (potentially out-of-order). Each subscriber's bounded channel
212    /// acts as the backpressure signal: when it fills, the adaptive ingestion controller
213    /// throttles fetch concurrency. The slowest subscriber gates ingestion for everyone.
214    ///
215    /// If a subscriber closes its channel, the ingestion service shuts down as well.
216    ///
217    /// If ingestion reaches the leading edge of the network, it will encounter checkpoints
218    /// that do not exist yet. These are retried on a fixed `retry_interval` until they become
219    /// available.
220    pub async fn run<R>(self, checkpoints: R) -> Result<Service>
221    where
222        R: std::ops::RangeBounds<u64> + Send + 'static,
223    {
224        let IngestionService {
225            config,
226            ingestion_client,
227            streaming_client,
228            gap_fill_client,
229            subscribers,
230        } = self;
231
232        if subscribers.is_empty() {
233            return Err(Error::NoSubscribers);
234        }
235
236        Ok(broadcaster(
237            checkpoints,
238            streaming_client,
239            config,
240            ingestion_client,
241            gap_fill_client,
242            subscribers,
243        ))
244    }
245}
246
247impl IngestionFactory {
248    /// Create a factory whose services fetch checkpoints as specified by `args`, reporting
249    /// metrics registered under `metrics_prefix` in `registry`.
250    pub(crate) async fn new(
251        args: ClientArgs,
252        config: IngestionConfig,
253        metrics_prefix: Option<&str>,
254        registry: &Registry,
255    ) -> Result<Self> {
256        let metrics = IngestionMetrics::new(metrics_prefix, registry);
257        let (ingestion_client, streaming_client, gap_fill_client) =
258            clients_from_args(args, &config, metrics).await?;
259        Ok(Self::with_clients(
260            ingestion_client,
261            streaming_client,
262            gap_fill_client,
263            config,
264        ))
265    }
266
267    /// Create a factory from pre-built clients, bypassing [ClientArgs]-driven construction.
268    /// Callers that supply their own [`IngestionClientTrait`] / [`CheckpointStreamingClient`]
269    /// implementations -- for example, when embedding the indexer in a fullnode that already has
270    /// checkpoint data on hand -- use this constructor instead of [`Self::new`].
271    ///
272    /// [`IngestionClientTrait`]: crate::ingestion::ingestion_client::IngestionClientTrait
273    pub(crate) fn with_clients(
274        ingestion_client: IngestionClient,
275        streaming_client: Option<ArcStreamingClient>,
276        gap_fill_client: Option<IngestionClient>,
277        config: IngestionConfig,
278    ) -> Self {
279        Self {
280            config,
281            ingestion_client,
282            streaming_client,
283            gap_fill_client,
284        }
285    }
286
287    /// Return the latest checkpoint number known to the underlying checkpoint source, preferably
288    /// via the streaming client, and failing that via the ingestion client.
289    pub(crate) async fn latest_checkpoint_number(&self) -> anyhow::Result<u64> {
290        latest_checkpoint_number(self.streaming_client.as_deref(), &self.ingestion_client).await
291    }
292
293    /// The ingestion client this factory's services fetch checkpoints with.
294    pub(crate) fn ingestion_client(&self) -> &IngestionClient {
295        &self.ingestion_client
296    }
297
298    /// The metrics that this factory's services report to.
299    pub(crate) fn metrics(&self) -> &Arc<IngestionMetrics> {
300        self.ingestion_client.metrics()
301    }
302
303    /// The ingestion configuration shared by every service this factory creates.
304    pub(crate) fn config(&self) -> &IngestionConfig {
305        &self.config
306    }
307
308    /// Create an ingestion service for `cohort` from the factory's clients. The cohort labels the
309    /// service's ingestion gauges so concurrently-running services don't overwrite each other.
310    pub(crate) fn create(&self, cohort: usize) -> IngestionService {
311        IngestionService::with_clients(
312            self.ingestion_client.for_cohort(cohort),
313            self.streaming_client.clone(),
314            self.gap_fill_client
315                .as_ref()
316                .map(|client| client.for_cohort(cohort)),
317            self.config.clone(),
318        )
319    }
320}
321
322impl Default for IngestionConfig {
323    fn default() -> Self {
324        Self {
325            ingest_concurrency: IngestConcurrencyConfig::Adaptive {
326                initial: 1,
327                min: 1,
328                max: 500,
329                dead_band: None,
330            },
331            retry_interval_ms: 200,
332            streaming_backoff_initial_batch_size: NonZeroUsize::new(10)
333                .expect("default initial streaming backoff is non-zero"), // 10 checkpoints, ~ 2 seconds
334            streaming_backoff_max_batch_size: 10000, // 10000 checkpoints, ~ 40 minutes
335            streaming_connection_timeout_ms: 5000,   // 5 seconds
336            streaming_statement_timeout_ms: 5000,    // 5 seconds
337            min_cohort_boundary: DEFAULT_MIN_COHORT_BOUNDARY,
338        }
339    }
340}
341
342/// Build the ingestion client, and the optional streaming and gap-fill clients, described by
343/// `args`.
344async fn clients_from_args(
345    args: ClientArgs,
346    config: &IngestionConfig,
347    metrics: Arc<IngestionMetrics>,
348) -> Result<(
349    IngestionClient,
350    Option<ArcStreamingClient>,
351    Option<IngestionClient>,
352)> {
353    let gap_fill_client = args
354        .streaming
355        .gap_fill_client(config.streaming_connection_timeout(), metrics.clone())
356        .await?;
357    let ingestion_client = IngestionClient::new(args.ingestion, metrics)?;
358    let streaming_client = args.streaming.streaming_url.map(|uri| {
359        Arc::new(GrpcStreamingClient::new(
360            uri,
361            config.streaming_connection_timeout(),
362            config.streaming_statement_timeout(),
363        )) as ArcStreamingClient
364    });
365    Ok((ingestion_client, streaming_client, gap_fill_client))
366}
367
368/// Return the latest checkpoint number known to the underlying checkpoint source, preferably
369/// via the streaming client, and failing that via the ingestion client.
370async fn latest_checkpoint_number<S>(
371    streaming_client: Option<&S>,
372    ingestion_client: &IngestionClient,
373) -> anyhow::Result<u64>
374where
375    S: CheckpointStreamingClient + Send + Sync + ?Sized,
376{
377    if let Some(streaming_client) = streaming_client {
378        match streaming_client.latest_checkpoint_number().await {
379            Ok(checkpoint_number) => return Ok(checkpoint_number),
380            Err(e) => {
381                warn!(
382                    operation = "latest_checkpoint_number",
383                    "Failed to get latest checkpoint number from streaming client: {e}"
384                );
385            }
386        }
387    }
388
389    let client = ingestion_client.clone();
390    let future = move || {
391        let client = client.clone();
392        async move {
393            client
394                .latest_checkpoint_number()
395                .await
396                .map_err(|e| backoff::Error::transient(Error::LatestCheckpointError(e)))
397        }
398    };
399
400    Ok(retry_transient_with_slow_monitor(
401        "latest_checkpoint_number",
402        future,
403        &ingestion_client
404            .metrics()
405            .ingested_latest_checkpoint_latency,
406    )
407    .await?)
408}
409
410#[cfg(test)]
411mod tests {
412    use std::sync::Mutex;
413
414    use axum::http::StatusCode;
415    use sui_futures::task::TaskGuard;
416    use url::Url;
417    use wiremock::MockServer;
418    use wiremock::Request;
419
420    use crate::ingestion::ingestion_client::tests::MockIngestionClient;
421    use crate::ingestion::store_client::tests::respond_with;
422    use crate::ingestion::store_client::tests::respond_with_chain_id;
423    use crate::ingestion::store_client::tests::status;
424    use crate::ingestion::streaming_client::test_utils::MockStreamingClient;
425    use crate::ingestion::test_utils::test_checkpoint_data;
426    use crate::metrics::IngestionMetrics;
427
428    use super::*;
429
430    const FALLBACK: u64 = 99;
431
432    fn mock_ingestion_client(latest_checkpoint: u64) -> IngestionClient {
433        let metrics = IngestionMetrics::new(None, &Registry::new());
434        let mock = MockIngestionClient {
435            latest_checkpoint,
436            ..Default::default()
437        };
438        IngestionClient::from_trait(Arc::new(mock), metrics)
439    }
440
441    async fn test_ingestion(uri: String, ingest_concurrency: usize) -> IngestionService {
442        let registry = Registry::new();
443        IngestionService::new(
444            ClientArgs {
445                ingestion: IngestionClientArgs {
446                    remote_store_url: Some(Url::parse(&uri).unwrap()),
447                    ..Default::default()
448                },
449                ..Default::default()
450            },
451            IngestionConfig {
452                ingest_concurrency: IngestConcurrencyConfig::Fixed {
453                    value: ingest_concurrency,
454                },
455                ..Default::default()
456            },
457            None,
458            &registry,
459        )
460        .await
461        .unwrap()
462    }
463
464    async fn test_subscriber(
465        stop_after: usize,
466        mut rx: mpsc::Receiver<Arc<CheckpointEnvelope>>,
467    ) -> TaskGuard<Vec<u64>> {
468        TaskGuard::new(tokio::spawn(async move {
469            let mut seqs = vec![];
470            for _ in 0..stop_after {
471                let Some(checkpoint_envelope) = rx.recv().await else {
472                    break;
473                };
474
475                seqs.push(checkpoint_envelope.checkpoint.summary.sequence_number);
476            }
477
478            seqs
479        }))
480    }
481
482    /// If the ingestion service has no subscribers, it will fail fast (before fetching any
483    /// checkpoints).
484    #[tokio::test]
485    async fn fail_on_no_subscribers() {
486        // The mock server will repeatedly return 404, so if the service does try to fetch a
487        // checkpoint, it will be stuck repeatedly retrying.
488        let server = MockServer::start().await;
489        respond_with(&server, status(StatusCode::NOT_FOUND)).await;
490
491        let ingestion_service = test_ingestion(server.uri(), 1).await;
492
493        let res = ingestion_service.run(0..).await;
494        assert!(matches!(res, Err(Error::NoSubscribers)));
495    }
496
497    /// The subscriber has no effective limit, and the mock server will always return checkpoint
498    /// information, but the ingestion service can still be stopped by shutting it down.
499    #[tokio::test]
500    async fn shutdown() {
501        let server = MockServer::start().await;
502        respond_with(
503            &server,
504            status(StatusCode::OK).set_body_bytes(test_checkpoint_data(42)),
505        )
506        .await;
507        respond_with_chain_id(&server).await;
508
509        let mut ingestion_service = test_ingestion(server.uri(), 1).await;
510
511        let rx = ingestion_service.subscribe_bounded(1);
512        let subscriber = test_subscriber(usize::MAX, rx).await;
513        let svc = ingestion_service.run(0..).await.unwrap();
514
515        svc.shutdown().await.unwrap();
516        subscriber.await.unwrap();
517    }
518
519    /// The subscriber will stop after receiving a single checkpoint, and this will trigger the
520    /// ingestion service to stop as well, even if there are more checkpoints to fetch.
521    #[tokio::test]
522    async fn shutdown_on_subscriber_drop() {
523        let server = MockServer::start().await;
524        respond_with(
525            &server,
526            status(StatusCode::OK).set_body_bytes(test_checkpoint_data(42)),
527        )
528        .await;
529        respond_with_chain_id(&server).await;
530
531        let mut ingestion_service = test_ingestion(server.uri(), 1).await;
532
533        let rx = ingestion_service.subscribe_bounded(1);
534        let subscriber = test_subscriber(1, rx).await;
535        let mut svc = ingestion_service.run(0..).await.unwrap();
536
537        drop(subscriber);
538        svc.join().await.unwrap();
539    }
540
541    /// The service will retry fetching a checkpoint that does not exist, in this test, the 4th
542    /// checkpoint will return 404 a couple of times, before eventually succeeding.
543    #[tokio::test]
544    async fn retry_on_not_found() {
545        let server = MockServer::start().await;
546        let times: Mutex<u64> = Mutex::new(0);
547        respond_with(&server, move |_: &Request| {
548            let mut times = times.lock().unwrap();
549            *times += 1;
550            match *times {
551                1..4 => status(StatusCode::OK).set_body_bytes(test_checkpoint_data(*times)),
552                4..6 => status(StatusCode::NOT_FOUND),
553                _ => status(StatusCode::OK).set_body_bytes(test_checkpoint_data(*times)),
554            }
555        })
556        .await;
557        respond_with_chain_id(&server).await;
558
559        let mut ingestion_service = test_ingestion(server.uri(), 1).await;
560
561        let rx = ingestion_service.subscribe_bounded(1);
562        let subscriber = test_subscriber(6, rx).await;
563        let _svc = ingestion_service.run(0..).await.unwrap();
564
565        let seqs = subscriber.await.unwrap();
566        assert_eq!(seqs, vec![0, 1, 2, 3, 6, 7]);
567    }
568
569    /// Similar to the previous test, but now it's a transient error that causes the retry.
570    #[tokio::test]
571    async fn retry_on_transient_error() {
572        let server = MockServer::start().await;
573        let times: Mutex<u64> = Mutex::new(0);
574        respond_with(&server, move |_: &Request| {
575            let mut times = times.lock().unwrap();
576            *times += 1;
577            match *times {
578                1..4 => status(StatusCode::OK).set_body_bytes(test_checkpoint_data(*times)),
579                4..6 => status(StatusCode::REQUEST_TIMEOUT),
580                _ => status(StatusCode::OK).set_body_bytes(test_checkpoint_data(*times)),
581            }
582        })
583        .await;
584        respond_with_chain_id(&server).await;
585
586        let mut ingestion_service = test_ingestion(server.uri(), 1).await;
587
588        let rx = ingestion_service.subscribe_bounded(1);
589        let subscriber = test_subscriber(6, rx).await;
590        let _svc = ingestion_service.run(0..).await.unwrap();
591
592        let seqs = subscriber.await.unwrap();
593        assert_eq!(seqs, vec![0, 1, 2, 3, 6, 7]);
594    }
595
596    /// One subscriber is going to stop processing checkpoints, so even though the service can keep
597    /// fetching checkpoints, it will stop short because of the slow receiver. Other subscribers
598    /// can keep processing checkpoints that were buffered for the slow one.
599    #[tokio::test]
600    async fn back_pressure_and_buffering() {
601        let server = MockServer::start().await;
602        let times: Mutex<u64> = Mutex::new(0);
603        respond_with(&server, move |_: &Request| {
604            let mut times = times.lock().unwrap();
605            *times += 1;
606            status(StatusCode::OK).set_body_bytes(test_checkpoint_data(*times))
607        })
608        .await;
609        respond_with_chain_id(&server).await;
610
611        let mut ingestion_service = test_ingestion(server.uri(), 1).await;
612        let size = 3;
613
614        // This subscriber will take its sweet time processing checkpoints.
615        let mut laggard = ingestion_service.subscribe_bounded(size);
616        async fn unblock(laggard: &mut mpsc::Receiver<Arc<CheckpointEnvelope>>) -> u64 {
617            let checkpoint_envelope = laggard.recv().await.unwrap();
618            checkpoint_envelope.checkpoint.summary.sequence_number
619        }
620
621        let rx = ingestion_service.subscribe_bounded(size);
622        let subscriber = test_subscriber(6, rx).await;
623        let _svc = ingestion_service.run(0..).await.unwrap();
624
625        // At this point, the service will have been able to pass 3 checkpoints to the non-lagging
626        // subscriber, while the laggard's buffer fills up. Now the laggard will pull two
627        // checkpoints, which will allow the rest of the pipeline to progress enough for the live
628        // subscriber to receive its quota. Checkpoint 0 is served by the chain_id mock.
629        assert_eq!(unblock(&mut laggard).await, 0);
630        assert_eq!(unblock(&mut laggard).await, 1);
631        assert_eq!(unblock(&mut laggard).await, 2);
632
633        let seqs = subscriber.await.unwrap();
634        assert_eq!(seqs, vec![0, 1, 2, 3, 4, 5]);
635    }
636
637    #[tokio::test]
638    async fn latest_checkpoint_number_no_streaming_client() {
639        let client = mock_ingestion_client(FALLBACK);
640        let result = latest_checkpoint_number(None::<&MockStreamingClient>, &client).await;
641        assert_eq!(result.unwrap(), FALLBACK);
642    }
643
644    #[tokio::test]
645    async fn latest_checkpoint_number_from_stream() {
646        let client = mock_ingestion_client(FALLBACK);
647        let streaming = Some(MockStreamingClient::new([42], None));
648        let result = latest_checkpoint_number(streaming.as_ref(), &client).await;
649        assert_eq!(result.unwrap(), 42);
650    }
651
652    #[tokio::test]
653    async fn latest_checkpoint_number_stream_error_falls_back() {
654        let client = mock_ingestion_client(FALLBACK);
655        let mut mock = MockStreamingClient::new(std::iter::empty::<u64>(), None);
656        mock.insert_error();
657        let result = latest_checkpoint_number(Some(&mock), &client).await;
658        assert_eq!(result.unwrap(), FALLBACK);
659    }
660
661    #[tokio::test]
662    async fn latest_checkpoint_number_empty_stream_falls_back() {
663        let client = mock_ingestion_client(FALLBACK);
664        let streaming = Some(MockStreamingClient::new(std::iter::empty::<u64>(), None));
665        let result = latest_checkpoint_number(streaming.as_ref(), &client).await;
666        assert_eq!(result.unwrap(), FALLBACK);
667    }
668
669    #[tokio::test]
670    async fn latest_checkpoint_number_connection_failure_falls_back() {
671        let client = mock_ingestion_client(FALLBACK);
672        let streaming = Some(
673            MockStreamingClient::new(std::iter::empty::<u64>(), None).fail_connection_times(1),
674        );
675        let result = latest_checkpoint_number(streaming.as_ref(), &client).await;
676        assert_eq!(result.unwrap(), FALLBACK);
677    }
678
679    #[test]
680    fn reject_zero_initial_streaming_backoff_batch_size() {
681        let mut config = serde_json::to_value(IngestionConfig::default()).unwrap();
682        config["streaming_backoff_initial_batch_size"] = serde_json::json!(0);
683
684        let error = serde_json::from_value::<IngestionConfig>(config).unwrap_err();
685        assert!(error.to_string().contains("nonzero"));
686    }
687
688    /// Every service an args-driven factory creates shares the factory's metrics; no additional
689    /// metrics are registered.
690    #[tokio::test]
691    async fn factory_from_args_shares_metrics() {
692        let registry = Registry::new();
693        let dir = tempfile::tempdir().unwrap();
694        let args = ClientArgs {
695            ingestion: IngestionClientArgs {
696                local_ingestion_path: Some(dir.path().to_owned()),
697                ..Default::default()
698            },
699            ..Default::default()
700        };
701
702        let factory = IngestionFactory::new(args, IngestionConfig::default(), None, &registry)
703            .await
704            .unwrap();
705
706        let first = factory.create(0);
707        let second = factory.create(1);
708        assert!(Arc::ptr_eq(
709            first.ingestion_client.metrics(),
710            factory.metrics()
711        ));
712        assert!(Arc::ptr_eq(
713            second.ingestion_client.metrics(),
714            factory.metrics()
715        ));
716        assert!(
717            registry
718                .gather()
719                .iter()
720                .all(|family| !family.name().contains("cohort"))
721        );
722    }
723
724    /// Services created by a client-driven factory all share the client's metrics handle, and
725    /// each gets its own clone of the streaming client. When the streaming client cannot serve
726    /// the latest checkpoint (the mock's shared queue is exhausted by the first probe), the
727    /// factory's probe falls back to the ingestion client.
728    #[tokio::test]
729    async fn factory_from_clients_shares_metrics() {
730        let registry = Registry::new();
731        let metrics = IngestionMetrics::new(None, &registry);
732        let client = IngestionClient::from_trait(
733            Arc::new(MockIngestionClient {
734                latest_checkpoint: FALLBACK,
735                ..Default::default()
736            }),
737            metrics.clone(),
738        );
739
740        let factory = IngestionFactory::with_clients(
741            client,
742            Some(Arc::new(MockStreamingClient::new([42], None))),
743            None,
744            IngestionConfig::default(),
745        );
746
747        assert_eq!(factory.latest_checkpoint_number().await.unwrap(), 42);
748
749        let first = factory.create(0);
750        let second = factory.create(1);
751        assert!(Arc::ptr_eq(first.ingestion_client.metrics(), &metrics));
752        assert!(Arc::ptr_eq(second.ingestion_client.metrics(), &metrics));
753        assert!(first.streaming_client.is_some());
754        assert!(second.streaming_client.is_some());
755
756        assert_eq!(factory.latest_checkpoint_number().await.unwrap(), FALLBACK);
757    }
758}
sui_indexer_alt_framework/ingestion/mod.rs

sui_indexer_alt_framework/ingestion/
mod.rs
