sui_indexer_alt_e2e_tests/

lib.rs

// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::collections::HashMap;
use std::fs;
use std::net::IpAddr;
use std::net::Ipv4Addr;
use std::net::SocketAddr;
use std::path::Path;
use std::time::Duration;

use anyhow::Context;
use anyhow::ensure;
use diesel::ExpressionMethods;
use diesel::OptionalExtension;
use diesel::QueryDsl;
use diesel_async::RunQueryDsl;
use prost::Message;
use reqwest::Client;
use serde_json::Value;
use serde_json::json;
use simulacrum::AdvanceEpochConfig;
use simulacrum::Simulacrum;
use sui_futures::service::Service;
use sui_indexer_alt::BootstrapGenesis;
use sui_indexer_alt::config::IndexerConfig;
use sui_indexer_alt::setup_indexer;
use sui_indexer_alt_consistent_api::proto::rpc::consistent::v1alpha::AvailableRangeRequest;
use sui_indexer_alt_consistent_api::proto::rpc::consistent::v1alpha::consistent_service_client::ConsistentServiceClient;
use sui_indexer_alt_consistent_store::args::RpcArgs as ConsistentArgs;
use sui_indexer_alt_consistent_store::args::TlsArgs as ConsistentTlsArgs;
use sui_indexer_alt_consistent_store::config::ServiceConfig as ConsistentConfig;
use sui_indexer_alt_consistent_store::start_service as start_consistent_store;
use sui_indexer_alt_framework::IndexerArgs;
use sui_indexer_alt_framework::ingestion::ClientArgs;
use sui_indexer_alt_framework::ingestion::ingestion_client::IngestionClientArgs;
use sui_indexer_alt_framework::pipeline::CommitterConfig;
use sui_indexer_alt_framework::postgres::schema::watermarks;
use sui_indexer_alt_graphql::RpcArgs as GraphQlArgs;
use sui_indexer_alt_graphql::args::SubscriptionArgs;
use sui_indexer_alt_graphql::config::RpcConfig as GraphQlConfig;
use sui_indexer_alt_graphql::start_rpc as start_graphql;
use sui_indexer_alt_reader::consistent_reader::ConsistentReaderArgs;
use sui_indexer_alt_reader::fullnode_client::FullnodeArgs;
use sui_indexer_alt_reader::kv_loader::KvArgs;
use sui_indexer_alt_reader::system_package_task::SystemPackageTaskArgs;
use sui_kv_rpc::KvRpcConfig;
use sui_kv_rpc::KvRpcServer;
use sui_kvstore::ALL_PIPELINE_NAMES;
use sui_kvstore::BigTableClient;
use sui_kvstore::BigTableIndexer;
use sui_kvstore::CHECKPOINTS_PIPELINE;
use sui_kvstore::IndexerConfig as BtIndexerConfig;
use sui_kvstore::IngestionConfig as BtIngestionConfig;
use sui_kvstore::KeyValueStoreReader;
use sui_kvstore::PipelineLayer;
use sui_kvstore::testing::BigTableEmulator;
use sui_kvstore::testing::INSTANCE_ID;
use sui_kvstore::testing::create_tables;
use sui_pg_db::Db;
use sui_pg_db::DbArgs;
use sui_pg_db::temp::TempDb;
use sui_pg_db::temp::get_available_port;
use sui_protocol_config::Chain;
use sui_rpc::field::FieldMask;
use sui_rpc::field::FieldMaskUtil;
use sui_rpc::merge::Merge;
use sui_rpc::proto::sui::rpc;
use sui_types::base_types::ObjectRef;
use sui_types::base_types::SuiAddress;
use sui_types::crypto::AccountKeyPair;
use sui_types::effects::TransactionEffects;
use sui_types::error::ExecutionError;
use sui_types::full_checkpoint_content::Checkpoint;
use sui_types::messages_checkpoint::VerifiedCheckpoint;
use sui_types::transaction::Transaction;
use tempfile::TempDir;
use tokio::time::error::Elapsed;
use tokio::time::interval;
use tokio::try_join;
use url::Url;

pub mod coin_registry;
pub mod find;
pub mod graphql;
pub mod move_helpers;
pub mod transaction;

/// A simulation of the network, accompanied by off-chain services (database, indexer, RPC),
/// connected by local data ingestion.
pub struct FullCluster {
    /// A simulation of the network, executing transactions and producing checkpoints.
    executor: Simulacrum,

    /// The off-chain services (database, indexer, RPC) that are ingesting data from the
    /// simulation.
    offchain: OffchainCluster,

    /// Temporary directory to store checkpoint information in, so that the indexer can pick it up.
    #[allow(unused)]
    temp_dir: TempDir,
}

/// A collection of the off-chain services (an indexer, a database, and a GraphQL server that reads
/// from that database), grouped together to simplify set-up and tear-down for tests. The included
/// RPC servers do not support transaction dry run and execution.
///
/// The database is temporary, and will be cleaned up when the cluster is dropped, and the RPCs are
/// set-up to listen on a random, available port, to avoid conflicts when multiple instances are
/// running concurrently in the same process.
pub struct OffchainCluster {
    /// The address the consistent store is listening on.
    consistent_listen_address: SocketAddr,

    /// The address the GraphQL server is listening on.
    graphql_listen_address: SocketAddr,

    /// The address the kv-rpc (LedgerService) server is listening on.
    kv_rpc_listen_address: SocketAddr,

    /// The address kv-rpc's second, unencrypted listener is listening on, when
    /// `OffchainClusterConfig::kv_rpc_plaintext_listener` is set.
    kv_rpc_plaintext_listen_address: Option<SocketAddr>,

    /// Read access to BigTable.
    bigtable_client: BigTableClient,

    /// Read access to the temporary database.
    db: Db,

    /// The pipelines that the indexer is populating.
    pipelines: Vec<&'static str>,

    /// Handles to all running services. Held on to so the services are not dropped (and therefore
    /// aborted) until the cluster is stopped.
    #[allow(unused)]
    services: Service,

    /// Handle to the BigTable emulator process.
    #[allow(unused)]
    bigtable_emulator: BigTableEmulator,

    /// Hold on to the database so it doesn't get dropped until the cluster is stopped.
    #[allow(unused)]
    database: TempDb,

    /// Hold on to the temporary directory where the consistent store writes its data, so it
    /// doesn't get cleaned up until the cluster is stopped.
    #[allow(unused)]
    dir: TempDir,
}

pub struct OffchainClusterConfig {
    pub indexer_args: IndexerArgs,
    pub consistent_indexer_args: IndexerArgs,
    pub fullnode_args: FullnodeArgs,
    pub indexer_config: IndexerConfig,
    pub consistent_config: ConsistentConfig,
    pub graphql_config: GraphQlConfig,
    pub bootstrap_genesis: Option<BootstrapGenesis>,
    pub kv_rpc_config: KvRpcConfig,
    /// Per-pipeline overrides (e.g. rate limits) for the BigTable archival indexer.
    pub bt_pipeline_layer: PipelineLayer,
    /// When set, kv-rpc also binds a second, unencrypted listener
    /// (`sui_kv_rpc::ServerConfig::plaintext_address`), reachable via
    /// `kv_rpc_plaintext_url`, for tests that exercise it directly.
    pub kv_rpc_plaintext_listener: bool,
}

impl FullCluster {
    /// Creates a cluster with a fresh executor where the off-chain services are set up with a
    /// default configuration.
    pub async fn new() -> anyhow::Result<Self> {
        Self::new_with_configs(
            Simulacrum::new(),
            OffchainClusterConfig::default(),
            &prometheus::Registry::new(),
        )
        .await
    }

    /// Creates a new cluster executing transactions using `executor`. The indexer is configured
    /// using `indexer_args` and `indexer_config`, and the GraphQL server is configured using
    /// `graphql_config`.
    pub async fn new_with_configs(
        mut executor: Simulacrum,
        offchain_cluster_config: OffchainClusterConfig,
        registry: &prometheus::Registry,
    ) -> anyhow::Result<Self> {
        let (client_args, temp_dir) = local_ingestion_client_args();
        executor.set_data_ingestion_path(temp_dir.path().to_owned());

        let offchain = OffchainCluster::new(client_args, offchain_cluster_config, registry)
            .await
            .context("Failed to create off-chain cluster")?;

        Ok(Self {
            executor,
            offchain,
            temp_dir,
        })
    }

    /// Return the reference gas price for the current epoch
    pub fn reference_gas_price(&self) -> u64 {
        self.executor.reference_gas_price()
    }

    /// Create a new account and credit it with `amount` gas units from a faucet account. Returns
    /// the account, its keypair, and a reference to the gas object it was funded with.
    pub fn funded_account(
        &mut self,
        amount: u64,
    ) -> anyhow::Result<(SuiAddress, AccountKeyPair, ObjectRef)> {
        self.executor.funded_account(amount)
    }

    /// Request gas from the faucet, sent to `address`. Return the object reference of the gas
    /// object that was sent.
    pub fn request_gas(
        &mut self,
        address: SuiAddress,
        amount: u64,
    ) -> anyhow::Result<TransactionEffects> {
        self.executor.request_gas(address, amount)
    }

    /// Execute a signed transaction, returning its effects.
    pub fn execute_transaction(
        &mut self,
        tx: Transaction,
    ) -> anyhow::Result<(TransactionEffects, Option<ExecutionError>)> {
        self.executor.execute_transaction(tx)
    }

    /// Execute a system transaction advancing the lock by the given `duration`.
    pub fn advance_clock(&mut self, duration: Duration) -> TransactionEffects {
        self.executor.advance_clock(duration)
    }

    /// Advance the executor into the next epoch. This executes an end-of-epoch transaction and
    /// creates the epoch's final checkpoint, but does not wait for the off-chain services to ingest
    /// it — follow with [`create_checkpoint`](Self::create_checkpoint) to sync.
    pub fn advance_epoch(&mut self) {
        self.executor.advance_epoch(AdvanceEpochConfig::default());
    }

    /// Create a new checkpoint containing the transactions executed since the last checkpoint that
    /// was created, and wait for the off-chain services to ingest it. Returns the checkpoint
    /// contents.
    pub async fn create_checkpoint(&mut self) -> VerifiedCheckpoint {
        let checkpoint = self.executor.create_checkpoint();
        let timeout = Duration::from_secs(100);
        let indexer = self
            .offchain
            .wait_for_indexer(checkpoint.sequence_number, timeout);
        let consistent_store = self
            .offchain
            .wait_for_consistent_store(checkpoint.sequence_number, timeout);
        let graphql = self
            .offchain
            .wait_for_graphql(checkpoint.sequence_number, timeout);
        let bigtable = self.offchain.wait_for_bigtable(
            &ALL_PIPELINE_NAMES,
            checkpoint.sequence_number,
            timeout,
        );
        let kv_rpc = self
            .offchain
            .wait_for_kv_rpc(checkpoint.sequence_number, timeout);

        try_join!(indexer, consistent_store, graphql, bigtable, kv_rpc)
            .expect("Timed out waiting for off-chain services");

        checkpoint
    }

    /// Unlike [`create_checkpoint`](Self::create_checkpoint), only waits for the base
    /// `checkpoints` BigTable pipeline to catch up — not the indexer, consistent store, GraphQL,
    /// or the list-index BigTable pipelines (`tx_seq_digest`, `transaction_bitmap_index`,
    /// `event_bitmap_index`). Used by tests that need to observe a checkpoint the base pipeline
    /// has indexed before the (typically throttled, via `bt_pipeline_layer`) list-index
    /// pipelines have processed it, without unrelated services' sync time giving the throttled
    /// pipelines room to catch up anyway.
    pub async fn create_checkpoint_before_list_apis_sync(&mut self) -> VerifiedCheckpoint {
        let checkpoint = self.executor.create_checkpoint();
        let timeout = Duration::from_secs(100);
        self.offchain
            .wait_for_bigtable(&[CHECKPOINTS_PIPELINE], checkpoint.sequence_number, timeout)
            .await
            .expect("Timed out waiting for the base checkpoints pipeline");

        checkpoint
    }

    /// Waits until every pipeline in `pipelines` has caught up to the given `checkpoint`, or the
    /// `timeout` is reached (an error).
    pub async fn wait_for_bigtable(
        &self,
        pipelines: &[&str],
        checkpoint: u64,
        timeout: Duration,
    ) -> Result<(), Elapsed> {
        self.offchain
            .wait_for_bigtable(pipelines, checkpoint, timeout)
            .await
    }

    /// The URL to talk to the database on.
    pub fn db_url(&self) -> Url {
        self.offchain.db_url()
    }

    /// The URL to send Consistent Store requests to.
    pub fn consistent_store_url(&self) -> Url {
        self.offchain.consistent_store_url()
    }

    /// The URL to send GraphQL requests to.
    pub fn graphql_url(&self) -> Url {
        self.offchain.graphql_url()
    }

    /// The URL to send kv-rpc (LedgerService) requests to.
    pub fn kv_rpc_url(&self) -> Url {
        self.offchain.kv_rpc_url()
    }

    /// The URL to send requests to kv-rpc's second, unencrypted listener, when
    /// `OffchainClusterConfig::kv_rpc_plaintext_listener` was set.
    pub fn kv_rpc_plaintext_url(&self) -> Option<Url> {
        self.offchain.kv_rpc_plaintext_url()
    }

    /// Returns the latest checkpoint that we have all data for in the database, according to the
    /// watermarks table. Returns `None` if any of the expected pipelines are missing data.
    pub async fn latest_checkpoint(&self) -> anyhow::Result<Option<u64>> {
        self.offchain.latest_checkpoint().await
    }

    /// Waits until the indexer has caught up to the given `checkpoint`, or the `timeout` is
    /// reached (an error).
    pub async fn wait_for_indexer(
        &self,
        checkpoint: u64,
        timeout: Duration,
    ) -> Result<(), Elapsed> {
        self.offchain.wait_for_indexer(checkpoint, timeout).await
    }

    /// Waits until the indexer's pruner has caught up to the given `checkpoint`, for the given
    /// `pipeline`, or the `timeout` is reached (an error).
    pub async fn wait_for_pruner(
        &self,
        pipeline: &str,
        checkpoint: u64,
        timeout: Duration,
    ) -> Result<(), Elapsed> {
        self.offchain
            .wait_for_pruner(pipeline, checkpoint, timeout)
            .await
    }

    /// Waits until GraphQL has caught up to the given `checkpoint`, or the `timeout` is
    /// reached (an error).
    pub async fn wait_for_graphql(
        &self,
        checkpoint: u64,
        timeout: Duration,
    ) -> Result<(), Elapsed> {
        self.offchain.wait_for_graphql(checkpoint, timeout).await
    }

    /// Waits until kv-rpc serves the given `checkpoint`, or the `timeout` is reached (an error).
    pub async fn wait_for_kv_rpc(&self, checkpoint: u64, timeout: Duration) -> Result<(), Elapsed> {
        self.offchain.wait_for_kv_rpc(checkpoint, timeout).await
    }
}

impl OffchainCluster {
    /// Construct a new off-chain cluster and spin up its constituent services.
    ///
    /// - `indexer_args`, `client_args`, and `indexer_config` control the indexer. In particular
    ///   `client_args` is used to configure the client that the indexer uses to fetch checkpoints.
    /// - `graphql_config` controls the GraphQL server.
    /// - `registry` is used to register metrics for the indexer and GraphQL servers.
    pub async fn new(
        client_args: ClientArgs,
        OffchainClusterConfig {
            indexer_args,
            consistent_indexer_args,
            fullnode_args,
            indexer_config,
            consistent_config,
            graphql_config,
            bootstrap_genesis,
            kv_rpc_config,
            bt_pipeline_layer,
            kv_rpc_plaintext_listener,
        }: OffchainClusterConfig,
        registry: &prometheus::Registry,
    ) -> anyhow::Result<Self> {
        let consistent_port = get_available_port();
        let consistent_listen_address =
            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), consistent_port);

        let graphql_port = get_available_port();
        let graphql_listen_address = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), graphql_port);

        let kv_rpc_port = get_available_port();
        let kv_rpc_address = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), kv_rpc_port);

        let kv_rpc_plaintext_address = kv_rpc_plaintext_listener.then(|| {
            let port = get_available_port();
            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port)
        });

        let database = TempDb::new().context("Failed to create database")?;
        let database_url = database.database().url();

        let dir = tempfile::tempdir().context("Failed to create temporary directory")?;
        let rocksdb_path = dir.path().join("rocksdb");

        let consistent_args = ConsistentArgs {
            rpc_listen_address: consistent_listen_address,
            tls: ConsistentTlsArgs::default(),
        };

        let graphql_args = GraphQlArgs {
            rpc_listen_address: graphql_listen_address,
            no_ide: true,
        };

        let db = Db::for_read(database_url.clone(), DbArgs::default())
            .await
            .context("Failed to connect to database")?;

        let indexer = setup_indexer(
            database_url.clone(),
            DbArgs::default(),
            indexer_args,
            client_args.clone(),
            indexer_config,
            bootstrap_genesis,
            registry,
        )
        .await
        .context("Failed to setup indexer")?;

        let pipelines: Vec<_> = indexer.pipelines().collect();
        let indexer = indexer.run().await.context("Failed to start indexer")?;

        let consistent_store = start_consistent_store(
            rocksdb_path,
            consistent_indexer_args,
            client_args.clone(),
            consistent_args,
            "0.0.0",
            consistent_config,
            registry,
        )
        .await
        .context("Failed to start Consistent Store")?;

        let consistent_reader_args = ConsistentReaderArgs {
            consistent_store_url: Some(
                Url::parse(&format!("http://{consistent_listen_address}")).unwrap(),
            ),
            ..Default::default()
        };

        let (bigtable_client, bigtable_emulator, archival_service) = start_archival(
            client_args.clone(),
            kv_rpc_address,
            kv_rpc_plaintext_address,
            kv_rpc_config,
            bt_pipeline_layer,
            registry,
        )
        .await?;

        let kv_args = KvArgs {
            ledger_grpc_url: Some(
                format!("http://{kv_rpc_address}")
                    .parse()
                    .expect("Failed to parse kv-rpc URI"),
            ),
            ..Default::default()
        };

        let graphql = start_graphql(
            Some(database_url.clone()),
            fullnode_args,
            DbArgs::default(),
            kv_args,
            consistent_reader_args,
            graphql_args,
            SystemPackageTaskArgs::default(),
            SubscriptionArgs::default(),
            "0.0.0",
            graphql_config,
            pipelines.iter().map(|p| p.to_string()).collect(),
            registry,
        )
        .await
        .context("Failed to start GraphQL server")?;

        let services = indexer
            .merge(consistent_store)
            .merge(graphql)
            .merge(archival_service);

        Ok(Self {
            consistent_listen_address,
            graphql_listen_address,
            kv_rpc_listen_address: kv_rpc_address,
            kv_rpc_plaintext_listen_address: kv_rpc_plaintext_address,
            bigtable_client,
            db,
            pipelines,
            services,
            bigtable_emulator,
            database,
            dir,
        })
    }

    /// The URL to talk to the database on.
    pub fn db_url(&self) -> Url {
        self.database.database().url().clone()
    }

    /// The URL to send Consistent Store requests to.
    pub fn consistent_store_url(&self) -> Url {
        Url::parse(&format!("http://{}/", self.consistent_listen_address))
            .expect("Failed to parse RPC URL")
    }

    /// The URL to send GraphQL requests to.
    pub fn graphql_url(&self) -> Url {
        Url::parse(&format!("http://{}/graphql", self.graphql_listen_address))
            .expect("Failed to parse RPC URL")
    }

    /// The URL to send kv-rpc (LedgerService) requests to.
    pub fn kv_rpc_url(&self) -> Url {
        Url::parse(&format!("http://{}/", self.kv_rpc_listen_address))
            .expect("Failed to parse RPC URL")
    }

    /// The URL to send requests to kv-rpc's second, unencrypted listener, when
    /// `OffchainClusterConfig::kv_rpc_plaintext_listener` was set.
    pub fn kv_rpc_plaintext_url(&self) -> Option<Url> {
        let address = self.kv_rpc_plaintext_listen_address?;
        Some(Url::parse(&format!("http://{address}/")).expect("Failed to parse RPC URL"))
    }

    /// Returns the latest checkpoint that we have all data for in the database, according to the
    /// watermarks table. Returns `None` if any of the expected pipelines are missing data.
    pub async fn latest_checkpoint(&self) -> anyhow::Result<Option<u64>> {
        use watermarks::dsl as w;

        let mut conn = self
            .db
            .connect()
            .await
            .context("Failed to connect to database")?;

        let latest: HashMap<String, i64> = w::watermarks
            .select((w::pipeline, w::checkpoint_hi_inclusive))
            .filter(w::pipeline.eq_any(&self.pipelines))
            .filter(w::reader_lo.le(w::checkpoint_hi_inclusive))
            .load(&mut conn)
            .await?
            .into_iter()
            .collect();

        if latest.len() != self.pipelines.len() {
            return Ok(None);
        }

        Ok(latest.into_values().min().map(|l| l as u64))
    }

    /// Returns the latest checkpoint that the pruner is willing to prune up to for the given
    /// `pipeline`.
    pub async fn latest_pruner_checkpoint(&self, pipeline: &str) -> anyhow::Result<Option<u64>> {
        use watermarks::dsl as w;

        let mut conn = self
            .db
            .connect()
            .await
            .context("Failed to connect to database")?;

        let latest: Option<i64> = w::watermarks
            .select(w::reader_lo)
            .filter(w::pipeline.eq(pipeline))
            .first(&mut conn)
            .await
            .optional()?;

        Ok(latest.map(|l| l as u64))
    }

    /// Returns the latest checkpoint that the consistent store is aware of.
    pub async fn latest_consistent_store_checkpoint(&self) -> anyhow::Result<u64> {
        ConsistentServiceClient::connect(self.consistent_store_url().to_string())
            .await
            .context("Failed to connect to Consistent Store")?
            .available_range(AvailableRangeRequest {})
            .await
            .context("Failed to fetch available range from Consistent Store")?
            .into_inner()
            .max_checkpoint
            .context("Consistent Store has not started yet")
    }

    /// Returns the latest checkpoint that the GraphQL service is aware of.
    pub async fn latest_graphql_checkpoint(&self) -> anyhow::Result<u64> {
        let query = json!({
            "query": "query { checkpoint { sequenceNumber } }"
        });

        let client = Client::new();
        let request = client.post(self.graphql_url()).json(&query);
        let response = request
            .send()
            .await
            .context("Request to GraphQL server failed")?;

        let body: Value = response
            .json()
            .await
            .context("Failed to parse GraphQL response")?;

        let sequence_number = body
            .pointer("/data/checkpoint/sequenceNumber")
            .context("Failed to find checkpoint sequence number in response")?;

        let sequence_number: i64 = serde_json::from_value(sequence_number.clone())
            .context("Failed to parse sequence number as i64")?;

        ensure!(sequence_number != i64::MAX, "Indexer has not started yet");

        Ok(sequence_number as u64)
    }

    /// Returns the latest epoch that the GraphQL service is aware of.
    pub async fn latest_graphql_epoch(&self) -> anyhow::Result<u64> {
        let query = json!({
            "query": "query { epoch { epochId } }"
        });

        let client = Client::new();
        let request = client.post(self.graphql_url()).json(&query);
        let response = request
            .send()
            .await
            .context("Request to GraphQL server failed")?;

        let body: Value = response
            .json()
            .await
            .context("Failed to parse GraphQL response")?;

        let epoch_id = body
            .pointer("/data/epoch/epochId")
            .context("Failed to find epochId in response")?;

        let epoch_id: i64 =
            serde_json::from_value(epoch_id.clone()).context("Failed to parse epochId as i64")?;

        ensure!(epoch_id != i64::MAX, "Indexer has not started yet");

        Ok(epoch_id as u64)
    }

    /// Waits until the indexer has caught up to the given `checkpoint`, or the `timeout` is
    /// reached (an error).
    pub async fn wait_for_indexer(
        &self,
        checkpoint: u64,
        timeout: Duration,
    ) -> Result<(), Elapsed> {
        tokio::time::timeout(timeout, async move {
            let mut interval = interval(Duration::from_millis(200));
            loop {
                interval.tick().await;
                if matches!(self.latest_checkpoint().await, Ok(Some(l)) if l >= checkpoint) {
                    break;
                }
            }
        })
        .await
    }

    /// Waits until the indexer's pruner has caught up to the given `checkpoint`, for the given
    /// `pipeline`, or the `timeout` is reached (an error).
    pub async fn wait_for_pruner(
        &self,
        pipeline: &str,
        checkpoint: u64,
        timeout: Duration,
    ) -> Result<(), Elapsed> {
        tokio::time::timeout(timeout, async move {
            let mut interval = interval(Duration::from_millis(200));
            loop {
                interval.tick().await;
                if matches!(self.latest_pruner_checkpoint(pipeline).await, Ok(Some(l)) if l >= checkpoint) {
                    break;
                }
            }
        }).await
    }

    /// Waits until the Consistent Store has caught up to the given `checkpoint`, or the `timeout`
    /// is reached (an error).
    pub async fn wait_for_consistent_store(
        &self,
        checkpoint: u64,
        timeout: Duration,
    ) -> Result<(), Elapsed> {
        tokio::time::timeout(timeout, async move {
            let mut interval = interval(Duration::from_millis(200));
            loop {
                interval.tick().await;
                if matches!(self.latest_consistent_store_checkpoint().await, Ok(l) if l >= checkpoint) {
                    break;
                }
            }
        })
        .await
    }

    /// Waits until GraphQL has caught up to the given `checkpoint`, or the `timeout` is reached
    /// (an error).
    pub async fn wait_for_graphql(
        &self,
        checkpoint: u64,
        timeout: Duration,
    ) -> Result<(), Elapsed> {
        tokio::time::timeout(timeout, async move {
            let mut interval = interval(Duration::from_millis(200));
            loop {
                interval.tick().await;
                if matches!(self.latest_graphql_checkpoint().await, Ok(l) if l >= checkpoint) {
                    break;
                }
            }
        })
        .await
    }

    /// Waits until kv-rpc's served height, which bounds every answer it gives, has reached the
    /// given `checkpoint`, or the `timeout` is reached (an error). The height comes from
    /// `GetServiceInfo`: a cached copy of the BigTable watermarks that kv-rpc refreshes in the
    /// background, so it can trail them.
    pub async fn wait_for_kv_rpc(&self, checkpoint: u64, timeout: Duration) -> Result<(), Elapsed> {
        let url = self.kv_rpc_url().to_string();
        tokio::time::timeout(timeout, async move {
            let mut interval = interval(Duration::from_millis(200));
            let mut client = None;
            loop {
                interval.tick().await;
                if client.is_none() {
                    client =
                        rpc::v2::ledger_service_client::LedgerServiceClient::connect(url.clone())
                            .await
                            .ok();
                }
                let Some(client) = client.as_mut() else {
                    continue;
                };
                if client
                    .get_service_info(rpc::v2::GetServiceInfoRequest::default())
                    .await
                    .is_ok_and(|response| {
                        response
                            .into_inner()
                            .checkpoint_height
                            .is_some_and(|cp| cp >= checkpoint)
                    })
                {
                    break;
                }
            }
        })
        .await
    }

    /// Waits until every pipeline in `pipelines` has caught up to the given `checkpoint`, or the
    /// `timeout` is reached (an error).
    pub async fn wait_for_bigtable(
        &self,
        pipelines: &[&str],
        checkpoint: u64,
        timeout: Duration,
    ) -> Result<(), Elapsed> {
        let mut client = self.bigtable_client.clone();
        tokio::time::timeout(timeout, async move {
            let mut interval = interval(Duration::from_millis(200));
            loop {
                interval.tick().await;
                if client
                    .get_watermark_for_pipelines(pipelines)
                    .await
                    .is_ok_and(|wm| {
                        wm.is_some_and(|wm| {
                            wm.checkpoint_hi_inclusive
                                .is_some_and(|cp| cp >= checkpoint)
                        })
                    })
                {
                    break;
                }
            }
        })
        .await
    }
}

impl Default for OffchainClusterConfig {
    fn default() -> Self {
        Self {
            indexer_args: Default::default(),
            consistent_indexer_args: Default::default(),
            fullnode_args: FullnodeArgs::default(),
            indexer_config: IndexerConfig::for_test(),
            consistent_config: ConsistentConfig::for_test(),
            graphql_config: Default::default(),
            bootstrap_genesis: None,
            kv_rpc_config: KvRpcConfig::default(),
            bt_pipeline_layer: PipelineLayer::default(),
            kv_rpc_plaintext_listener: false,
        }
    }
}

/// Returns ClientArgs that use a temporary local ingestion path and the TempDir of that path.
pub fn local_ingestion_client_args() -> (ClientArgs, TempDir) {
    let temp_dir = tempfile::tempdir()
        .context("Failed to create data ingestion path")
        .unwrap();
    let client_args = ClientArgs {
        ingestion: IngestionClientArgs {
            local_ingestion_path: Some(temp_dir.path().to_owned()),
            ..Default::default()
        },
        ..Default::default()
    };
    (client_args, temp_dir)
}

/// Writes a checkpoint file to the given path.
pub async fn write_checkpoint(path: &Path, checkpoint: Checkpoint) -> anyhow::Result<()> {
    let sequence_number = checkpoint.summary.sequence_number;

    let mask = FieldMask::from_paths([
        rpc::v2::Checkpoint::path_builder().sequence_number(),
        rpc::v2::Checkpoint::path_builder().summary().bcs().value(),
        rpc::v2::Checkpoint::path_builder().signature().finish(),
        rpc::v2::Checkpoint::path_builder().contents().bcs().value(),
        rpc::v2::Checkpoint::path_builder()
            .transactions()
            .transaction()
            .bcs()
            .value(),
        rpc::v2::Checkpoint::path_builder()
            .transactions()
            .effects()
            .bcs()
            .value(),
        rpc::v2::Checkpoint::path_builder()
            .transactions()
            .effects()
            .unchanged_loaded_runtime_objects()
            .finish(),
        rpc::v2::Checkpoint::path_builder()
            .transactions()
            .events()
            .bcs()
            .value(),
        rpc::v2::Checkpoint::path_builder()
            .objects()
            .objects()
            .bcs()
            .value(),
    ]);

    let proto_checkpoint = rpc::v2::Checkpoint::merge_from(&checkpoint, &mask.into());
    let proto_bytes = proto_checkpoint.encode_to_vec();
    let compressed = zstd::encode_all(&proto_bytes[..], 3)?;

    let file_name = format!("{}.binpb.zst", sequence_number);
    let file_path = path.join(file_name);
    fs::write(file_path, compressed)?;
    Ok(())
}

/// Start the archival stack: BigTable emulator, BigTable indexer, and sui-kv-rpc.
async fn start_archival(
    client_args: ClientArgs,
    kv_rpc_address: SocketAddr,
    kv_rpc_plaintext_address: Option<SocketAddr>,
    kv_rpc_config: KvRpcConfig,
    bt_pipeline_layer: PipelineLayer,
    registry: &prometheus::Registry,
) -> anyhow::Result<(BigTableClient, BigTableEmulator, Service)> {
    let emulator = tokio::task::spawn_blocking(BigTableEmulator::start)
        .await
        .context("spawn_blocking panicked")?
        .context("Failed to start BigTable emulator")?;

    create_tables(emulator.host(), INSTANCE_ID)
        .await
        .context("Failed to create BigTable tables")?;

    let bigtable_client =
        BigTableClient::new_local(emulator.host().to_string(), INSTANCE_ID.to_string())
            .await
            .context("Failed to create BigTable client")?;

    let indexer_client =
        BigTableClient::new_local(emulator.host().to_string(), INSTANCE_ID.to_string())
            .await
            .context("Failed to create BigTable client for indexer")?;
    let bt_indexer = BigTableIndexer::new(
        indexer_client,
        IndexerArgs::default(),
        client_args,
        BtIngestionConfig::default(),
        CommitterConfig::default(),
        BtIndexerConfig::default(),
        bt_pipeline_layer,
        &[],
        Chain::Unknown,
        registry,
    )
    .await
    .context("Failed to create BigTable indexer")?;

    // Use the BigTable wrapper, not the raw framework indexer, so bitmap
    // committer background tasks are supervised for the duration of the test.
    let bt_indexer_service = bt_indexer
        .run()
        .await
        .context("Failed to start BigTable indexer")?;

    let kv_rpc_server = KvRpcServer::new_local_with_config(
        emulator.host().to_string(),
        INSTANCE_ID.to_string(),
        None,
        kv_rpc_config.ledger_history(),
        kv_rpc_config.request_bigtable_concurrency(),
        kv_rpc_config.monotonic_read_wait_timeout(),
        kv_rpc_config.stages(),
        kv_rpc_config.watermark_source(),
    )
    .await
    .context("Failed to create KvRpcServer")?;
    let kv_rpc_service = kv_rpc_server
        .start_service(
            kv_rpc_address,
            sui_kv_rpc::ServerConfig {
                plaintext_address: kv_rpc_plaintext_address,
                ..Default::default()
            },
        )
        .await
        .context("Failed to start kv-rpc server")?;

    let service = bt_indexer_service.merge(kv_rpc_service);
    Ok((bigtable_client, emulator, service))
}