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

sui_indexer_alt_framework/ingestion/
mod.rs
