Skip to main content

sui_indexer_alt_framework/ingestion/
ingestion_client.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::future::Future;
5use std::path::PathBuf;
6use std::sync::Arc;
7use std::time::Duration;
8
9use async_trait::async_trait;
10use backoff::Error as BE;
11use backoff::ExponentialBackoff;
12use backoff::backoff::Constant;
13use clap::ArgGroup;
14use object_store::ClientOptions;
15use object_store::ObjectStore;
16use object_store::aws::AmazonS3Builder;
17use object_store::azure::MicrosoftAzureBuilder;
18use object_store::gcp::GoogleCloudStorageBuilder;
19use object_store::http::HttpBuilder;
20use object_store::local::LocalFileSystem;
21use prometheus::Histogram;
22use reqwest::header::HeaderMap;
23use reqwest::header::HeaderName;
24use reqwest::header::HeaderValue;
25use sui_futures::future::with_slow_future_monitor;
26use sui_http::middleware::callback::CallbackLayer;
27use sui_rpc::Client;
28use sui_rpc::client::HeadersInterceptor;
29use sui_types::digests::ChainIdentifier;
30use tokio::sync::OnceCell;
31use tracing::debug;
32use tracing::warn;
33use url::Url;
34
35use crate::ingestion::Error as IE;
36use crate::ingestion::MAX_GRPC_MESSAGE_SIZE_BYTES;
37use crate::ingestion::Result as IngestionResult;
38use crate::ingestion::byte_count::ByteCountMakeCallbackHandler;
39use crate::ingestion::decode;
40use crate::ingestion::store_client::StoreIngestionClient;
41use crate::metrics::CohortMetrics;
42use crate::metrics::IngestionMetrics;
43use crate::types::full_checkpoint_content::Checkpoint;
44
45/// Wait at most this long between retries for transient errors.
46const MAX_TRANSIENT_RETRY_INTERVAL: Duration = Duration::from_secs(60);
47
48/// Threshold for logging warnings about slow HTTP operations during checkpoint fetching.
49///
50/// Operations that take longer than this duration will trigger a warning log, but will
51/// continue executing without being canceled. This helps identify network issues or
52/// slow remote stores without interrupting the ingestion process.
53const SLOW_OPERATION_WARNING_THRESHOLD: Duration = Duration::from_secs(60);
54
55/// The `cohort` metric label for the base client, before it has been rebound to a specific cohort via
56/// [`IngestionClient::for_cohort`]. It keeps checkpoint fetches made before the indexer starts (e.g.
57/// the factory's tip probe) out of any real cohort's series.
58const DEFAULT_COHORT_LABEL: &str = "";
59
60#[async_trait]
61pub trait IngestionClientTrait: Send + Sync {
62    async fn chain_id(&self) -> anyhow::Result<ChainIdentifier>;
63
64    async fn checkpoint(&self, checkpoint: u64) -> CheckpointResult;
65
66    async fn latest_checkpoint_number(&self) -> anyhow::Result<u64>;
67}
68
69#[derive(clap::Args, Clone, Debug)]
70#[command(group(ArgGroup::new("source").required(true).multiple(false)))]
71pub struct IngestionClientArgs {
72    /// Remote Store to fetch checkpoints from over HTTP.
73    #[arg(long, group = "source")]
74    pub remote_store_url: Option<Url>,
75
76    /// Fetch checkpoints from AWS S3. Provide the bucket name or endpoint-and-bucket.
77    /// (env: AWS_ENDPOINT, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_DEFAULT_REGION)
78    #[arg(long, group = "source")]
79    pub remote_store_s3: Option<String>,
80
81    /// Fetch checkpoints from Google Cloud Storage. Provide the bucket name.
82    /// (env: GOOGLE_SERVICE_ACCOUNT_PATH)
83    #[arg(long, group = "source")]
84    pub remote_store_gcs: Option<String>,
85
86    /// Fetch checkpoints from Azure Blob Storage. Provide the container name.
87    /// (env: AZURE_STORAGE_ACCOUNT_NAME, AZURE_STORAGE_ACCESS_KEY)
88    #[arg(long, group = "source")]
89    pub remote_store_azure: Option<String>,
90
91    /// Default header to include in remote store requests, as `<name>:<value>`.
92    /// Can be provided multiple times.
93    #[arg(long = "remote-store-header", value_parser = parse_remote_store_header)]
94    pub remote_store_headers: Vec<(HeaderName, HeaderValue)>,
95
96    /// Path to the local ingestion directory.
97    #[arg(long, group = "source")]
98    pub local_ingestion_path: Option<PathBuf>,
99
100    /// Sui fullnode gRPC url to fetch checkpoints from.
101    #[arg(long, group = "source")]
102    pub rpc_api_url: Option<Url>,
103
104    /// Optional username for the gRPC service.
105    #[arg(long, env, requires = "rpc_api_url")]
106    pub rpc_username: Option<String>,
107
108    /// Optional password for the gRPC service.
109    #[arg(long, env, requires = "rpc_api_url")]
110    pub rpc_password: Option<String>,
111
112    /// How long to wait for a checkpoint file to be downloaded (milliseconds). Set to 0 to disable
113    /// the timeout.
114    #[arg(long, default_value_t = Self::default().checkpoint_timeout_ms)]
115    pub checkpoint_timeout_ms: u64,
116
117    /// How long to wait while establishing a connection to the checkpoint store (milliseconds).
118    /// Set to 0 to disable the timeout.
119    #[arg(long, default_value_t = Self::default().checkpoint_connection_timeout_ms)]
120    pub checkpoint_connection_timeout_ms: u64,
121}
122
123impl Default for IngestionClientArgs {
124    fn default() -> Self {
125        Self {
126            remote_store_url: None,
127            remote_store_s3: None,
128            remote_store_gcs: None,
129            remote_store_azure: None,
130            remote_store_headers: vec![],
131            local_ingestion_path: None,
132            rpc_api_url: None,
133            rpc_username: None,
134            rpc_password: None,
135            checkpoint_timeout_ms: 120_000,
136            checkpoint_connection_timeout_ms: 120_000,
137        }
138    }
139}
140
141impl IngestionClientArgs {
142    fn client_options(&self) -> ClientOptions {
143        let mut options = ClientOptions::default();
144
145        options = if self.checkpoint_timeout_ms == 0 {
146            options.with_timeout_disabled()
147        } else {
148            let timeout = Duration::from_millis(self.checkpoint_timeout_ms);
149            options.with_timeout(timeout)
150        };
151
152        options = if self.checkpoint_connection_timeout_ms == 0 {
153            options.with_connect_timeout_disabled()
154        } else {
155            let timeout = Duration::from_millis(self.checkpoint_connection_timeout_ms);
156            options.with_connect_timeout(timeout)
157        };
158
159        options = if !self.remote_store_headers.is_empty() {
160            let mut headers = HeaderMap::new();
161            for (name, value) in &self.remote_store_headers {
162                headers.append(name.clone(), value.clone());
163            }
164
165            options.with_default_headers(headers)
166        } else {
167            options
168        };
169
170        options
171    }
172}
173
174#[derive(thiserror::Error, Debug)]
175pub enum CheckpointError {
176    #[error("Checkpoint not found")]
177    NotFound,
178    #[error("Failed to fetch checkpoint: {0}")]
179    Fetch(#[from] anyhow::Error),
180    #[error("Failed to decode checkpoint: {0}")]
181    Decode(#[from] decode::Error),
182}
183
184pub type CheckpointResult = Result<Checkpoint, CheckpointError>;
185
186#[derive(Clone)]
187pub struct IngestionClient {
188    client: Arc<dyn IngestionClientTrait>,
189    /// The metrics registry this client's cohort view is bound against, and the handle its peers
190    /// reuse (see [`Self::metrics`]). Also owns the global `total_ingested_bytes` counter wired into
191    /// the checkpoint source.
192    metrics: Arc<IngestionMetrics>,
193    /// This client's cohort-labeled metric handles, rebound per cohort by [`Self::for_cohort`].
194    cohort_metrics: Arc<CohortMetrics>,
195    /// Shared so cloning an uninitialized client does not create an independent cache.
196    chain_id: Arc<OnceCell<ChainIdentifier>>,
197}
198
199#[derive(Clone, Debug)]
200pub struct CheckpointEnvelope {
201    pub checkpoint: Arc<Checkpoint>,
202    pub chain_id: ChainIdentifier,
203}
204
205impl IngestionClient {
206    /// Construct a new ingestion client. Its source is determined by `args`.
207    pub fn new(args: IngestionClientArgs, metrics: Arc<IngestionMetrics>) -> IngestionResult<Self> {
208        // TODO: Support stacking multiple ingestion clients for redundancy/failover.
209        let retry = super::store_client::retry_config();
210        let client = if let Some(url) = args.remote_store_url.as_ref() {
211            let store = HttpBuilder::new()
212                .with_url(url.to_string())
213                .with_client_options(args.client_options().with_allow_http(true))
214                .with_retry(retry)
215                .build()
216                .map(Arc::new)?;
217            IngestionClient::with_store(store, metrics.clone())?
218        } else if let Some(bucket) = args.remote_store_s3.as_ref() {
219            let store = AmazonS3Builder::from_env()
220                .with_client_options(args.client_options())
221                .with_retry(retry)
222                .with_imdsv1_fallback()
223                .with_bucket_name(bucket)
224                .build()
225                .map(Arc::new)?;
226            IngestionClient::with_store(store, metrics.clone())?
227        } else if let Some(bucket) = args.remote_store_gcs.as_ref() {
228            let store = GoogleCloudStorageBuilder::from_env()
229                .with_client_options(args.client_options())
230                .with_retry(retry)
231                .with_bucket_name(bucket)
232                .build()
233                .map(Arc::new)?;
234            IngestionClient::with_store(store, metrics.clone())?
235        } else if let Some(container) = args.remote_store_azure.as_ref() {
236            let store = MicrosoftAzureBuilder::from_env()
237                .with_client_options(args.client_options())
238                .with_retry(retry)
239                .with_container_name(container)
240                .build()
241                .map(Arc::new)?;
242            IngestionClient::with_store(store, metrics.clone())?
243        } else if let Some(path) = args.local_ingestion_path.as_ref() {
244            let store = LocalFileSystem::new_with_prefix(path).map(Arc::new)?;
245            IngestionClient::with_store(store, metrics.clone())?
246        } else if let Some(rpc_api_url) = args.rpc_api_url.as_ref() {
247            IngestionClient::with_grpc(
248                rpc_api_url.clone(),
249                args.rpc_username,
250                args.rpc_password,
251                metrics.clone(),
252            )?
253        } else {
254            panic!(
255                "One of remote_store_url, remote_store_s3, remote_store_gcs, remote_store_azure, \
256                local_ingestion_path or rpc_api_url must be provided"
257            );
258        };
259
260        Ok(client)
261    }
262
263    /// An ingestion client that fetches checkpoints from a remote object store.
264    pub fn with_store(
265        store: Arc<dyn ObjectStore>,
266        metrics: Arc<IngestionMetrics>,
267    ) -> IngestionResult<Self> {
268        let client = Arc::new(StoreIngestionClient::new(
269            store,
270            Some(metrics.total_ingested_bytes.clone()),
271        ));
272        Ok(Self::from_trait(client, metrics))
273    }
274
275    /// An ingestion client that fetches checkpoints from a fullnode, over gRPC.
276    pub fn with_grpc(
277        url: Url,
278        username: Option<String>,
279        password: Option<String>,
280        metrics: Arc<IngestionMetrics>,
281    ) -> IngestionResult<Self> {
282        let byte_count_layer = tower::ServiceBuilder::new()
283            .layer(CallbackLayer::new(ByteCountMakeCallbackHandler::new(
284                metrics.total_ingested_bytes.clone(),
285            )))
286            // Rebox the callback middleware's wrapped request body back into the
287            // body type that the underlying tonic channel is monomorphic on.
288            .map_request(|request: http::Request<_>| request.map(tonic::body::Body::new));
289        let client = Client::new(url.to_string())?
290            .with_max_decoding_message_size(MAX_GRPC_MESSAGE_SIZE_BYTES)
291            .request_layer(byte_count_layer);
292        let client = if let Some(username) = username {
293            let mut headers = HeadersInterceptor::new();
294            headers.basic_auth(username, password);
295            client.with_headers(headers)
296        } else {
297            client
298        };
299        Ok(Self::from_trait(Arc::new(client), metrics))
300    }
301
302    /// The metrics handle this client reports against. Callers constructing peer services (e.g. an
303    /// [`IngestionService`]) against the same client should reuse this Arc rather than building a
304    /// second [`IngestionMetrics`] from the same registry, which would double-register the metric
305    /// vectors.
306    ///
307    /// [`IngestionService`]: crate::ingestion::IngestionService
308    pub fn metrics(&self) -> &Arc<IngestionMetrics> {
309        &self.metrics
310    }
311
312    /// This client's cohort-labeled metric view, bound to the cohort it was created for.
313    pub(crate) fn cohort_metrics(&self) -> &Arc<CohortMetrics> {
314        &self.cohort_metrics
315    }
316
317    /// Wrap an arbitrary [`IngestionClientTrait`] implementation in an [`IngestionClient`]. Use
318    /// this when the source of checkpoints is not one of the built-in remote object stores or gRPC
319    /// endpoints — for example, when embedding the indexer in a fullnode that already has
320    /// checkpoint data on hand.
321    pub fn from_trait(
322        client: Arc<dyn IngestionClientTrait>,
323        metrics: Arc<IngestionMetrics>,
324    ) -> Self {
325        let cohort_metrics = CohortMetrics::new(&metrics, DEFAULT_COHORT_LABEL);
326        IngestionClient {
327            client,
328            metrics,
329            cohort_metrics,
330            chain_id: Arc::new(OnceCell::new()),
331        }
332    }
333
334    /// A clone of this client that reports its cohort-labeled metrics under `cohort`, so services
335    /// minted for different cohorts don't overwrite each other. The checkpoint source and chain-id
336    /// cache are shared; only the cohort metric view is rebound.
337    pub(crate) fn for_cohort(&self, cohort: usize) -> Self {
338        let mut client = self.clone();
339        client.cohort_metrics = CohortMetrics::new(&self.metrics, &cohort.to_string());
340        client
341    }
342
343    /// Fetch checkpoint data by sequence number.
344    ///
345    /// This function behaves like `IngestionClient::fetch`, but will repeatedly retry the fetch if
346    /// the checkpoint is not found, on a constant back-off. The time between fetches is controlled
347    /// by the `retry_interval` parameter.
348    pub async fn wait_for(
349        &self,
350        checkpoint: u64,
351        retry_interval: Duration,
352    ) -> IngestionResult<CheckpointEnvelope> {
353        let backoff = Constant::new(retry_interval);
354        let fetch = || async move {
355            use backoff::Error as BE;
356            self.checkpoint(checkpoint).await.map_err(|e| match e {
357                IE::NotFound(checkpoint) => {
358                    debug!(checkpoint, "Checkpoint not found, retrying...");
359                    self.cohort_metrics.total_ingested_not_found_retries.inc();
360                    BE::transient(e)
361                }
362                e => BE::permanent(e),
363            })
364        };
365
366        backoff::future::retry(backoff, fetch).await
367    }
368
369    /// Fetch checkpoint data by sequence number.
370    ///
371    /// Repeatedly retries transient errors with an exponential backoff (up to
372    /// `MAX_TRANSIENT_RETRY_INTERVAL`). Transient errors are defined by the client
373    /// implementation that returns a [CheckpointError::Fetch] or [CheckpointError::Decode] error
374    /// variant.
375    ///
376    /// The function will immediately return if the checkpoint is not found.
377    pub async fn checkpoint(&self, cp_sequence_number: u64) -> IngestionResult<CheckpointEnvelope> {
378        let client = self.client.clone();
379        let checkpoint_data_fut = retry_transient_with_slow_monitor(
380            "checkpoint",
381            move || {
382                let client = client.clone();
383                async move {
384                    client
385                        .checkpoint(cp_sequence_number)
386                        .await
387                        .map_err(|err| match err {
388                            // Not found errors are marked as permanent here, but retried in
389                            // `wait_for` in case the checkpoint becomes available in the future.
390                            CheckpointError::NotFound => {
391                                BE::permanent(IE::NotFound(cp_sequence_number))
392                            }
393                            // Retry fetch and decode errors in case the root cause is in the
394                            // upstream checkpoint data source. If the upstream checkpoint data
395                            // source is corrected, then the indexer will automatically recover
396                            // the next time the read is attempted.
397                            CheckpointError::Fetch(e) => self.cohort_metrics.inc_retry(
398                                cp_sequence_number,
399                                "fetch",
400                                IE::FetchError(cp_sequence_number, e),
401                            ),
402                            CheckpointError::Decode(e) => self.cohort_metrics.inc_retry(
403                                cp_sequence_number,
404                                e.reason(),
405                                IE::DecodeError(cp_sequence_number, e.into()),
406                            ),
407                        })
408                }
409            },
410            &self.cohort_metrics.ingested_checkpoint_latency,
411        );
412
413        let client = self.client.clone();
414        let chain_id_fut = self.chain_id.get_or_try_init(|| {
415            retry_transient_with_slow_monitor(
416                "chain_id",
417                move || {
418                    let client = client.clone();
419                    async move {
420                        client
421                            .chain_id()
422                            .await
423                            .map_err(|e| BE::transient(IE::ChainIdError(cp_sequence_number, e)))
424                    }
425                },
426                &self.cohort_metrics.ingested_chain_id_latency,
427            )
428        });
429
430        let (checkpoint, chain_id) = tokio::try_join!(checkpoint_data_fut, chain_id_fut)?;
431
432        self.cohort_metrics
433            .checkpoint_lag
434            .report_lag(cp_sequence_number, checkpoint.summary.timestamp_ms);
435
436        self.cohort_metrics.total_ingested_checkpoints.inc();
437
438        self.cohort_metrics
439            .total_ingested_transactions
440            .inc_by(checkpoint.transactions.len() as u64);
441
442        self.cohort_metrics.total_ingested_events.inc_by(
443            checkpoint
444                .transactions
445                .iter()
446                .map(|tx| tx.events.as_ref().map_or(0, |evs| evs.data.len()) as u64)
447                .sum(),
448        );
449
450        self.cohort_metrics
451            .total_ingested_objects
452            .inc_by(checkpoint.object_set.len() as u64);
453
454        Ok(CheckpointEnvelope {
455            checkpoint: Arc::new(checkpoint),
456            chain_id: *chain_id,
457        })
458    }
459
460    pub async fn latest_checkpoint_number(&self) -> anyhow::Result<u64> {
461        self.client.latest_checkpoint_number().await
462    }
463}
464
465/// Keep backing off until we are waiting for the max interval, but don't give up.
466pub(crate) fn transient_backoff() -> ExponentialBackoff {
467    ExponentialBackoff {
468        max_interval: MAX_TRANSIENT_RETRY_INTERVAL,
469        max_elapsed_time: None,
470        ..Default::default()
471    }
472}
473
474/// Retry a fallible async operation with exponential backoff and slow-operation monitoring.
475/// Records the total time (including retries) in the provided latency histogram.
476pub(crate) async fn retry_transient_with_slow_monitor<F, Fut, T>(
477    operation: &str,
478    make_future: F,
479    latency: &Histogram,
480) -> IngestionResult<T>
481where
482    F: Fn() -> Fut,
483    Fut: Future<Output = Result<T, backoff::Error<IE>>>,
484{
485    let request = || {
486        let fut = make_future();
487        async move {
488            with_slow_future_monitor(fut, SLOW_OPERATION_WARNING_THRESHOLD, || {
489                warn!(
490                    operation,
491                    threshold_ms = SLOW_OPERATION_WARNING_THRESHOLD.as_millis(),
492                    "Slow operation detected"
493                );
494            })
495            .await
496        }
497    };
498
499    let guard = latency.start_timer();
500    let data = backoff::future::retry(transient_backoff(), request).await?;
501    let elapsed = guard.stop_and_record();
502
503    debug!(
504        operation,
505        elapsed_ms = elapsed * 1000.0,
506        "Fetched operation"
507    );
508
509    Ok(data)
510}
511
512fn parse_remote_store_header(header: &str) -> Result<(HeaderName, HeaderValue), String> {
513    let (name, value) = header
514        .split_once(':')
515        .ok_or_else(|| "remote store header must be in `<name>:<value>` format".to_string())?;
516
517    let name = HeaderName::from_bytes(name.as_bytes())
518        .map_err(|err| format!("invalid remote store header name `{name}`: {err}"))?;
519    let value = HeaderValue::from_str(value)
520        .map_err(|err| format!("invalid remote store header value for `{name}`: {err}"))?;
521
522    Ok((name, value))
523}
524
525#[cfg(test)]
526pub(crate) mod tests {
527    use std::sync::Arc;
528    use std::sync::atomic::AtomicUsize;
529    use std::sync::atomic::Ordering;
530    use std::time::Duration;
531
532    use clap::Parser;
533    use clap::error::ErrorKind;
534    use dashmap::DashMap;
535    use prometheus::Registry;
536    use sui_types::digests::CheckpointDigest;
537    use sui_types::event::Event;
538    use sui_types::test_checkpoint_data_builder::TestCheckpointBuilder;
539
540    use crate::ingestion::decode;
541    use crate::ingestion::test_utils::test_checkpoint_data;
542
543    use super::*;
544
545    fn test_checkpoint(seq: u64) -> Checkpoint {
546        let bytes = test_checkpoint_data(seq);
547        decode::checkpoint(&bytes).unwrap()
548    }
549
550    /// Build a checkpoint with one transaction containing one event and one created object.
551    fn test_checkpoint_with_data(seq: u64) -> Checkpoint {
552        TestCheckpointBuilder::new(seq)
553            .start_transaction(0)
554            .create_owned_object(0)
555            .with_events(vec![Event::random_for_testing()])
556            .finish_transaction()
557            .build_checkpoint()
558    }
559
560    #[derive(Debug, Parser)]
561    struct TestArgs {
562        #[clap(flatten)]
563        ingestion: IngestionClientArgs,
564    }
565
566    /// Mock implementation of IngestionClientTrait for testing.
567    ///
568    /// - `checkpoints`: pre-inserted checkpoints returned by `checkpoint()`. Sequences not
569    ///   in the map return `CheckpointError::NotFound`.
570    /// - `not_found_failures` / `fetch_failures` / `decode_failures`: number of times to
571    ///   return the corresponding error for a given sequence number before succeeding.
572    /// - `latest_checkpoint`: value returned by `latest_checkpoint_number()`.
573    #[derive(Default)]
574    pub(crate) struct MockIngestionClient {
575        pub checkpoints: DashMap<u64, Checkpoint>,
576        pub not_found_failures: DashMap<u64, usize>,
577        pub fetch_failures: DashMap<u64, usize>,
578        pub decode_failures: DashMap<u64, usize>,
579        pub latest_checkpoint: u64,
580        pub chain_id_calls: AtomicUsize,
581    }
582
583    impl MockIngestionClient {
584        pub(crate) fn mock_chain_id() -> ChainIdentifier {
585            CheckpointDigest::new([1; 32]).into()
586        }
587
588        /// Populate `checkpoints` with synthetic test checkpoints for the given sequence
589        /// numbers.
590        pub(crate) fn insert_checkpoints(&self, range: impl IntoIterator<Item = u64>) {
591            for seq in range {
592                self.checkpoints.insert(seq, test_checkpoint(seq));
593            }
594        }
595    }
596
597    #[async_trait]
598    impl IngestionClientTrait for MockIngestionClient {
599        async fn chain_id(&self) -> anyhow::Result<ChainIdentifier> {
600            self.chain_id_calls.fetch_add(1, Ordering::Relaxed);
601            Ok(Self::mock_chain_id())
602        }
603
604        async fn checkpoint(&self, checkpoint: u64) -> CheckpointResult {
605            if let Some(mut remaining) = self.not_found_failures.get_mut(&checkpoint)
606                && *remaining > 0
607            {
608                *remaining -= 1;
609                return Err(CheckpointError::NotFound);
610            }
611
612            if let Some(mut remaining) = self.fetch_failures.get_mut(&checkpoint)
613                && *remaining > 0
614            {
615                *remaining -= 1;
616                return Err(CheckpointError::Fetch(anyhow::anyhow!("Mock fetch error")));
617            }
618
619            if let Some(mut remaining) = self.decode_failures.get_mut(&checkpoint)
620                && *remaining > 0
621            {
622                *remaining -= 1;
623                return Err(CheckpointError::Decode(decode::Error::Deserialization(
624                    prost::DecodeError::new("Mock deserialization error"),
625                )));
626            }
627
628            self.checkpoints
629                .get(&checkpoint)
630                .as_deref()
631                .cloned()
632                .ok_or(CheckpointError::NotFound)
633        }
634
635        async fn latest_checkpoint_number(&self) -> anyhow::Result<u64> {
636            Ok(self.latest_checkpoint)
637        }
638    }
639
640    fn setup_test() -> (IngestionClient, Arc<MockIngestionClient>) {
641        let registry = Registry::new_custom(Some("test".to_string()), None).unwrap();
642        let metrics = IngestionMetrics::new(None, &registry);
643        let mock_client = Arc::new(MockIngestionClient::default());
644        // Bind to cohort 0, matching production where the broadcaster always runs under a cohort's
645        // client rather than the unlabeled base client.
646        let client = IngestionClient::from_trait(mock_client.clone(), metrics).for_cohort(0);
647        (client, mock_client)
648    }
649
650    #[test]
651    fn test_args_multiple_ingestion_sources_are_rejected() {
652        let err = TestArgs::try_parse_from([
653            "cmd",
654            "--remote-store-url",
655            "https://example.com",
656            "--rpc-api-url",
657            "http://localhost:8080",
658        ])
659        .unwrap_err();
660
661        assert_eq!(err.kind(), ErrorKind::ArgumentConflict);
662    }
663
664    #[test]
665    fn test_args_optional_credentials() {
666        let args = TestArgs::try_parse_from([
667            "cmd",
668            "--rpc-api-url",
669            "http://localhost:8080",
670            "--rpc-username",
671            "alice",
672            "--rpc-password",
673            "secret",
674        ])
675        .unwrap();
676
677        assert_eq!(args.ingestion.rpc_username.as_deref(), Some("alice"));
678        assert_eq!(args.ingestion.rpc_password.as_deref(), Some("secret"));
679        assert_eq!(
680            args.ingestion.rpc_api_url,
681            Some(Url::parse("http://localhost:8080").unwrap())
682        );
683    }
684
685    #[test]
686    fn test_args_credentials_require_rpc_url() {
687        let err = TestArgs::try_parse_from([
688            "cmd",
689            "--rpc-username",
690            "alice",
691            "--rpc-password",
692            "secret",
693        ])
694        .unwrap_err();
695
696        assert_eq!(err.kind(), ErrorKind::MissingRequiredArgument);
697    }
698
699    #[test]
700    fn test_args_remote_store_headers() {
701        let args = TestArgs::try_parse_from([
702            "cmd",
703            "--remote-store-gcs",
704            "bucket",
705            "--remote-store-header",
706            "x-goog-user-project:my-project",
707            "--remote-store-header",
708            "authorization:Bearer abc:def",
709        ])
710        .unwrap();
711
712        assert_eq!(args.ingestion.remote_store_headers.len(), 2);
713        assert_eq!(
714            args.ingestion.remote_store_headers[0].0,
715            HeaderName::from_static("x-goog-user-project")
716        );
717        assert_eq!(
718            args.ingestion.remote_store_headers[0].1,
719            HeaderValue::from_static("my-project")
720        );
721        assert_eq!(
722            args.ingestion.remote_store_headers[1].0,
723            HeaderName::from_static("authorization")
724        );
725        assert_eq!(
726            args.ingestion.remote_store_headers[1].1,
727            HeaderValue::from_static("Bearer abc:def")
728        );
729    }
730
731    #[test]
732    fn test_args_remote_store_header_requires_delimiter() {
733        let err = TestArgs::try_parse_from([
734            "cmd",
735            "--remote-store-gcs",
736            "bucket",
737            "--remote-store-header",
738            "x-goog-user-project",
739        ])
740        .unwrap_err();
741
742        assert_eq!(err.kind(), ErrorKind::ValueValidation);
743    }
744
745    #[test]
746    fn test_args_remote_store_header_rejects_invalid_name() {
747        let err = TestArgs::try_parse_from([
748            "cmd",
749            "--remote-store-gcs",
750            "bucket",
751            "--remote-store-header",
752            "bad name:value",
753        ])
754        .unwrap_err();
755
756        assert_eq!(err.kind(), ErrorKind::ValueValidation);
757    }
758
759    #[test]
760    fn test_args_remote_store_header_rejects_invalid_value() {
761        let err = TestArgs::try_parse_from([
762            "cmd",
763            "--remote-store-gcs",
764            "bucket",
765            "--remote-store-header",
766            "x-test:bad\nvalue",
767        ])
768        .unwrap_err();
769
770        assert_eq!(err.kind(), ErrorKind::ValueValidation);
771    }
772
773    #[tokio::test]
774    async fn test_checkpoint_checkpoint_success() {
775        let (client, mock) = setup_test();
776
777        mock.checkpoints.insert(1, test_checkpoint_with_data(1));
778
779        let result = client.checkpoint(1).await.unwrap();
780        assert_eq!(result.checkpoint.summary.sequence_number(), &1);
781        assert_eq!(result.chain_id, MockIngestionClient::mock_chain_id());
782        assert_eq!(
783            client
784                .metrics
785                .total_ingested_checkpoints
786                .with_label_values(&["0"])
787                .get(),
788            1
789        );
790        assert_eq!(
791            client
792                .metrics
793                .total_ingested_transactions
794                .with_label_values(&["0"])
795                .get(),
796            1
797        );
798        assert_eq!(
799            client
800                .metrics
801                .total_ingested_events
802                .with_label_values(&["0"])
803                .get(),
804            1
805        );
806        // 1 created object + 2 gas object versions (input + output)
807        assert_eq!(
808            client
809                .metrics
810                .total_ingested_objects
811                .with_label_values(&["0"])
812                .get(),
813            3
814        );
815    }
816
817    #[tokio::test]
818    async fn test_clones_share_chain_id_cache() {
819        let (client, mock_client) = setup_test();
820        mock_client.insert_checkpoints(1..=2);
821
822        let first = client.clone();
823        let second = client.clone();
824        let (first, second) = tokio::join!(first.checkpoint(1), second.checkpoint(2));
825
826        assert_eq!(
827            first.unwrap().chain_id,
828            MockIngestionClient::mock_chain_id()
829        );
830        assert_eq!(
831            second.unwrap().chain_id,
832            MockIngestionClient::mock_chain_id()
833        );
834        assert_eq!(mock_client.chain_id_calls.load(Ordering::Relaxed), 1);
835    }
836
837    #[tokio::test]
838    async fn test_checkpoint_not_found() {
839        let (client, _) = setup_test();
840
841        // Try to fetch non-existent checkpoint
842        let result = client.checkpoint(1).await;
843        assert!(matches!(result, Err(IE::NotFound(1))));
844        assert_eq!(
845            client
846                .metrics
847                .total_ingested_checkpoints
848                .with_label_values(&["0"])
849                .get(),
850            0
851        );
852        assert_eq!(
853            client
854                .metrics
855                .total_ingested_transactions
856                .with_label_values(&["0"])
857                .get(),
858            0
859        );
860        assert_eq!(
861            client
862                .metrics
863                .total_ingested_events
864                .with_label_values(&["0"])
865                .get(),
866            0
867        );
868        assert_eq!(
869            client
870                .metrics
871                .total_ingested_objects
872                .with_label_values(&["0"])
873                .get(),
874            0
875        );
876    }
877
878    #[tokio::test]
879    async fn test_checkpoint_fetch_error_with_retry() {
880        let (client, mock) = setup_test();
881
882        mock.checkpoints.insert(1, test_checkpoint(1));
883        mock.fetch_failures.insert(1, 2);
884
885        // Fetch and verify it succeeds after retries
886        let result = client.checkpoint(1).await.unwrap();
887        assert_eq!(*result.checkpoint.summary.sequence_number(), 1);
888        assert_eq!(result.chain_id, MockIngestionClient::mock_chain_id());
889
890        // Verify that exactly 2 retries were recorded
891        let retries = client
892            .metrics
893            .total_ingested_transient_retries
894            .with_label_values(&["fetch", "0"])
895            .get();
896        assert_eq!(retries, 2);
897        assert_eq!(
898            client
899                .metrics
900                .total_ingested_checkpoints
901                .with_label_values(&["0"])
902                .get(),
903            1
904        );
905        assert_eq!(
906            client
907                .metrics
908                .total_ingested_transactions
909                .with_label_values(&["0"])
910                .get(),
911            0
912        );
913        assert_eq!(
914            client
915                .metrics
916                .total_ingested_events
917                .with_label_values(&["0"])
918                .get(),
919            0
920        );
921        assert_eq!(
922            client
923                .metrics
924                .total_ingested_objects
925                .with_label_values(&["0"])
926                .get(),
927            0
928        );
929    }
930
931    #[tokio::test]
932    async fn test_checkpoint_decode_error_with_retry() {
933        let (client, mock) = setup_test();
934
935        mock.checkpoints.insert(1, test_checkpoint(1));
936        mock.decode_failures.insert(1, 2);
937
938        // Fetch and verify it succeeds after retries
939        let result = client.checkpoint(1).await.unwrap();
940        assert_eq!(*result.checkpoint.summary.sequence_number(), 1);
941        assert_eq!(result.chain_id, MockIngestionClient::mock_chain_id());
942
943        // Verify that exactly 2 retries were recorded
944        let retries = client
945            .metrics
946            .total_ingested_transient_retries
947            .with_label_values(&["deserialization", "0"])
948            .get();
949        assert_eq!(retries, 2);
950        assert_eq!(
951            client
952                .metrics
953                .total_ingested_checkpoints
954                .with_label_values(&["0"])
955                .get(),
956            1
957        );
958        assert_eq!(
959            client
960                .metrics
961                .total_ingested_transactions
962                .with_label_values(&["0"])
963                .get(),
964            0
965        );
966        assert_eq!(
967            client
968                .metrics
969                .total_ingested_events
970                .with_label_values(&["0"])
971                .get(),
972            0
973        );
974        assert_eq!(
975            client
976                .metrics
977                .total_ingested_objects
978                .with_label_values(&["0"])
979                .get(),
980            0
981        );
982    }
983
984    #[tokio::test]
985    async fn test_wait_for_checkpoint_with_retry() {
986        let (client, mock) = setup_test();
987
988        mock.checkpoints.insert(1, test_checkpoint(1));
989        mock.not_found_failures.insert(1, 1);
990
991        // Wait for checkpoint with short retry interval
992        let result = client.wait_for(1, Duration::from_millis(50)).await.unwrap();
993        assert_eq!(result.checkpoint.summary.sequence_number(), &1);
994        assert_eq!(result.chain_id, MockIngestionClient::mock_chain_id());
995
996        // Verify that exactly 1 retry was recorded
997        let retries = client
998            .metrics
999            .total_ingested_not_found_retries
1000            .with_label_values(&["0"])
1001            .get();
1002        assert_eq!(retries, 1);
1003        assert_eq!(
1004            client
1005                .metrics
1006                .total_ingested_checkpoints
1007                .with_label_values(&["0"])
1008                .get(),
1009            1
1010        );
1011        assert_eq!(
1012            client
1013                .metrics
1014                .total_ingested_transactions
1015                .with_label_values(&["0"])
1016                .get(),
1017            0
1018        );
1019        assert_eq!(
1020            client
1021                .metrics
1022                .total_ingested_events
1023                .with_label_values(&["0"])
1024                .get(),
1025            0
1026        );
1027        assert_eq!(
1028            client
1029                .metrics
1030                .total_ingested_objects
1031                .with_label_values(&["0"])
1032                .get(),
1033            0
1034        );
1035    }
1036
1037    #[tokio::test]
1038    async fn test_wait_for_checkpoint_instant() {
1039        let (client, mock) = setup_test();
1040
1041        mock.checkpoints.insert(1, test_checkpoint(1));
1042
1043        let result = client.wait_for(1, Duration::from_millis(50)).await.unwrap();
1044        assert_eq!(result.checkpoint.summary.sequence_number(), &1);
1045        assert_eq!(result.chain_id, MockIngestionClient::mock_chain_id());
1046        assert_eq!(
1047            client
1048                .metrics
1049                .total_ingested_checkpoints
1050                .with_label_values(&["0"])
1051                .get(),
1052            1
1053        );
1054        assert_eq!(
1055            client
1056                .metrics
1057                .total_ingested_transactions
1058                .with_label_values(&["0"])
1059                .get(),
1060            0
1061        );
1062        assert_eq!(
1063            client
1064                .metrics
1065                .total_ingested_events
1066                .with_label_values(&["0"])
1067                .get(),
1068            0
1069        );
1070        assert_eq!(
1071            client
1072                .metrics
1073                .total_ingested_objects
1074                .with_label_values(&["0"])
1075                .get(),
1076            0
1077        );
1078    }
1079}