[state sync] migrate to new bucket format

Author: phoenix-o

Cargo.lock

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use std::path::PathBuf;
+use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -14,8 +15,11 @@ use clap::ArgGroup;
 use object_store::ClientOptions;
 use object_store::ObjectStore;
 use object_store::aws::AmazonS3Builder;
+use object_store::aws::AmazonS3ConfigKey;
+use object_store::azure::AzureConfigKey;
 use object_store::azure::MicrosoftAzureBuilder;
 use object_store::gcp::GoogleCloudStorageBuilder;
+use object_store::gcp::GoogleConfigKey;
 use object_store::http::HttpBuilder;
 use object_store::local::LocalFileSystem;
 use sui_futures::future::with_slow_future_monitor;
@@ -102,6 +106,10 @@ pub struct IngestionClientArgs {
     /// Set to 0 to disable the timeout.
     #[arg(long, default_value_t = Self::default().checkpoint_connection_timeout_ms)]
     pub checkpoint_connection_timeout_ms: u64,
+
+    /// Extra provider-specific settings for the remote store client(S3, GCS, and Azure backends).
+    #[arg(skip)]
+    pub remote_store_extra_settings: Vec<(String, String)>,
 }
 
 impl Default for IngestionClientArgs {
@@ -117,6 +125,7 @@ impl Default for IngestionClientArgs {
             rpc_password: None,
             checkpoint_timeout_ms: 120_000,
             checkpoint_connection_timeout_ms: 120_000,
+            remote_store_extra_settings: vec![],
         }
     }
 }
@@ -196,30 +205,37 @@ impl IngestionClient {
                 .map(Arc::new)?;
             IngestionClient::with_store(store, metrics.clone())?
         } else if let Some(bucket) = args.remote_store_s3.as_ref() {
-            let store = AmazonS3Builder::from_env()
+            let mut builder = AmazonS3Builder::from_env()
                 .with_client_options(args.client_options())
                 .with_retry(retry)
-                .with_imdsv1_fallback()
-                .with_bucket_name(bucket)
-                .build()
-                .map(Arc::new)?;
-            IngestionClient::with_store(store, metrics.clone())?
+                .with_imdsv1_fallback();
+            builder = if Url::parse(bucket).is_ok() {
+                builder.with_url(bucket)
+            } else {
+                builder.with_bucket_name(bucket)
+            };
+            for (key, value) in &args.remote_store_extra_settings {
+                builder = builder.with_config(AmazonS3ConfigKey::from_str(key)?, value.clone());
+            }
+            IngestionClient::with_store(Arc::new(builder.build()?), metrics.clone())?
         } else if let Some(bucket) = args.remote_store_gcs.as_ref() {
-            let store = GoogleCloudStorageBuilder::from_env()
+            let mut builder = GoogleCloudStorageBuilder::from_env()
                 .with_client_options(args.client_options())
                 .with_retry(retry)
-                .with_bucket_name(bucket)
-                .build()
-                .map(Arc::new)?;
-            IngestionClient::with_store(store, metrics.clone())?
+                .with_bucket_name(bucket);
+            for (key, value) in &args.remote_store_extra_settings {
+                builder = builder.with_config(GoogleConfigKey::from_str(key)?, value.clone());
+            }
+            IngestionClient::with_store(Arc::new(builder.build()?), metrics.clone())?
         } else if let Some(container) = args.remote_store_azure.as_ref() {
-            let store = MicrosoftAzureBuilder::from_env()
+            let mut builder = MicrosoftAzureBuilder::from_env()
                 .with_client_options(args.client_options())
                 .with_retry(retry)
-                .with_container_name(container)
-                .build()
-                .map(Arc::new)?;
-            IngestionClient::with_store(store, metrics.clone())?
+                .with_container_name(container);
+            for (key, value) in &args.remote_store_extra_settings {
+                builder = builder.with_config(AzureConfigKey::from_str(key)?, value.clone());
+            }
+            IngestionClient::with_store(Arc::new(builder.build()?), metrics.clone())?
         } else if let Some(path) = args.local_ingestion_path.as_ref() {
             let store = LocalFileSystem::new_with_prefix(path).map(Arc::new)?;
             IngestionClient::with_store(store, metrics.clone())?

crates/sui-indexer-alt-framework/src/ingestion/ingestion_client.rs

@@ -35,7 +35,8 @@ sui-config.workspace = true
 sui-swarm-config.workspace = true
 sui-simulator.workspace = true
 sui-protocol-config.workspace = true
-sui-data-ingestion-core.workspace = true
+sui-indexer-alt-framework.workspace = true
+url.workspace = true
 
 arc-swap.workspace = true
 bcs.workspace = true

crates/sui-network/Cargo.toml

@@ -7,12 +7,22 @@ use prometheus::{
     register_int_counter_with_registry, register_int_gauge_with_registry,
 };
 use std::sync::Arc;
+use sui_indexer_alt_framework::metrics::IngestionMetrics;
 use sui_types::messages_checkpoint::CheckpointSequenceNumber;
 use tap::Pipe;
 
 #[derive(Clone)]
 pub(super) struct Metrics(Option<Arc<Inner>>);
 
+impl Metrics {
+    pub fn ingestion_metrics(&self) -> Arc<IngestionMetrics> {
+        self.0
+            .as_ref()
+            .map(|inner| inner.ingestion_metrics.clone())
+            .unwrap_or_else(|| IngestionMetrics::new(None, &Registry::new()))
+    }
+}
+
 impl std::fmt::Debug for Metrics {
     fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
         fmt.debug_struct("Metrics").finish()
@@ -87,6 +97,7 @@ struct Inner {
     checkpoint_summary_age: Histogram,
     // TODO: delete once users are migrated to non-Mysten histogram.
     checkpoint_summary_age_ms: MystenHistogram,
+    ingestion_metrics: Arc<IngestionMetrics>,
 }
 
 impl Inner {
@@ -136,6 +147,7 @@ impl Inner {
                 "Age of checkpoints summaries when they arrive and are verified.",
                 registry,
             ),
+            ingestion_metrics: IngestionMetrics::new(None, registry),
         }
         .pipe(Arc::new)
     }

crates/sui-network/src/state_sync/metrics.rs

@@ -84,7 +84,7 @@ mod tests;
 mod worker;
 
 use self::{metrics::Metrics, server::CheckpointContentsDownloadLimitLayer};
-use crate::state_sync::worker::StateSyncWorker;
+use crate::state_sync::worker::process_archive_checkpoint;
 pub use builder::{Builder, UnstartedStateSync};
 pub use generated::{
     state_sync_client::StateSyncClient,
@@ -93,7 +93,9 @@ pub use generated::{
 pub use server::GetCheckpointAvailabilityResponse;
 pub use server::GetCheckpointSummaryRequest;
 use sui_config::node::ArchiveReaderConfig;
-use sui_data_ingestion_core::{ReaderOptions, setup_single_workflow_with_options};
+use sui_indexer_alt_framework::ingestion::ingestion_client::{
+    IngestionClient, IngestionClientArgs,
+};
 use sui_storage::verify_checkpoint;
 
 /// A handle to the StateSync subsystem.
@@ -1576,29 +1578,44 @@ async fn sync_checkpoint_contents_from_archive_iteration<S>(
             warn!("{} can't be used as an archival fallback", ingestion_url);
             return;
         }
-        let reader_options = ReaderOptions {
-            batch_size: archive_config.download_concurrency.into(),
-            upper_limit: Some(end),
-            ..Default::default()
-        };
-        let Ok((executor, _exit_sender)) = setup_single_workflow_with_options(
-            StateSyncWorker(store, metrics),
-            ingestion_url.clone(),
-            archive_config.remote_store_options.clone(),
-            start,
-            1,
-            Some(reader_options),
-        )
-        .await
-        else {
-            return;
+        let url = ingestion_url
+            .parse::<url::Url>()
+            .expect("archival ingestion url must be valid");
+        let args = if url.host_str().unwrap_or_default().starts_with("s3") {
+            IngestionClientArgs {
+                remote_store_s3: Some(ingestion_url.clone()),
+                remote_store_extra_settings: archive_config.remote_store_options.clone(),
+                ..Default::default()
+            }
+        } else {
+            IngestionClientArgs {
+                remote_store_url: Some(url),
+                remote_store_extra_settings: archive_config.remote_store_options.clone(),
+                ..Default::default()
+            }
         };
-        match executor.await {
-            Ok(_) => info!(
-                "State sync from archive is complete. Checkpoints downloaded = {:?}",
-                end - start
-            ),
-            Err(err) => warn!("State sync from archive failed with error: {:?}", err),
+        let client = IngestionClient::new(args, metrics.ingestion_metrics())
+            .expect("failed to create ingestion client for state sync");
+        let mut checkpoint_stream = futures::stream::iter(start..=end)
+            .map(|seq| {
+                let client = client.clone();
+                async move { (seq, client.checkpoint(seq).await) }
+            })
+            .buffered(archive_config.download_concurrency.get());
+
+        while let Some((seq, result)) = checkpoint_stream.next().await {
+            let Some(checkpoint) = result
+                .inspect_err(|e| {
+                    warn!("State sync from archive failed fetching checkpoint {seq}: {e}")
+                })
+                .ok()
+            else {
+                return;
+            };
+            if let Err(e) = process_archive_checkpoint(&store, &checkpoint.checkpoint, &metrics) {
+                warn!("State sync from archive failed processing checkpoint {seq}: {e}");
+                return;
+            }
         }
     }
 }

crates/sui-network/src/state_sync/mod.rs

@@ -2,42 +2,43 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::state_sync::metrics::Metrics;
-use anemo::async_trait;
 use anyhow::{Context, anyhow};
-use sui_data_ingestion_core::Worker;
 use sui_storage::verify_checkpoint;
-use sui_types::full_checkpoint_content::CheckpointData;
+use sui_types::base_types::ExecutionData;
+use sui_types::full_checkpoint_content::Checkpoint;
 use sui_types::messages_checkpoint::CertifiedCheckpointSummary;
 use sui_types::messages_checkpoint::VerifiedCheckpoint;
 use sui_types::messages_checkpoint::VerifiedCheckpointContents;
 use sui_types::messages_checkpoint::VersionedFullCheckpointContents;
 use sui_types::storage::WriteStore;
+use sui_types::transaction::Transaction;
 
-pub(crate) struct StateSyncWorker<S>(pub(crate) S, pub(crate) Metrics);
-
-#[async_trait]
-impl<S: WriteStore + Clone + Send + Sync + 'static> Worker for StateSyncWorker<S> {
-    type Result = ();
-
-    async fn process_checkpoint(&self, checkpoint: &CheckpointData) -> anyhow::Result<()> {
-        let verified_checkpoint = get_or_insert_verified_checkpoint(
-            &self.0,
-            checkpoint.checkpoint_summary.clone(),
-            true,
-        )?;
-        let full_contents = VersionedFullCheckpointContents::from_contents_and_execution_data(
-            checkpoint.checkpoint_contents.clone(),
-            checkpoint.transactions.iter().map(|t| t.execution_data()),
-        );
-        full_contents.verify_digests(verified_checkpoint.content_digest)?;
-        let verified_contents = VerifiedCheckpointContents::new_unchecked(full_contents);
-        self.0
-            .insert_checkpoint_contents(&verified_checkpoint, verified_contents)?;
-        self.0
-            .update_highest_synced_checkpoint(&verified_checkpoint)?;
-        self.1.update_checkpoints_synced_from_archive();
-        Ok(())
-    }
+pub(crate) fn process_archive_checkpoint<S>(
+    store: &S,
+    checkpoint: &Checkpoint,
+    metrics: &Metrics,
+) -> anyhow::Result<()>
+where
+    S: WriteStore + Clone,
+{
+    let verified_checkpoint =
+        get_or_insert_verified_checkpoint(store, checkpoint.summary.clone(), true)?;
+    let full_contents = VersionedFullCheckpointContents::from_contents_and_execution_data(
+        checkpoint.contents.clone(),
+        checkpoint.transactions.iter().map(|t| ExecutionData {
+            transaction: Transaction::from_generic_sig_data(
+                t.transaction.clone(),
+                t.signatures.clone(),
+            ),
+            effects: t.effects.clone(),
+        }),
+    );
+    full_contents.verify_digests(verified_checkpoint.content_digest)?;
+    let verified_contents = VerifiedCheckpointContents::new_unchecked(full_contents);
+    store.insert_checkpoint_contents(&verified_checkpoint, verified_contents)?;
+    store.update_highest_synced_checkpoint(&verified_checkpoint)?;
+    metrics.update_checkpoints_synced_from_archive();
+    Ok(())
 }
 
 pub fn get_or_insert_verified_checkpoint<S>(