Skip to main content

sui_indexer_alt_e2e_tests/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::HashMap;
5use std::fs;
6use std::net::IpAddr;
7use std::net::Ipv4Addr;
8use std::net::SocketAddr;
9use std::path::Path;
10use std::time::Duration;
11
12use anyhow::Context;
13use anyhow::ensure;
14use diesel::ExpressionMethods;
15use diesel::OptionalExtension;
16use diesel::QueryDsl;
17use diesel_async::RunQueryDsl;
18use prost::Message;
19use reqwest::Client;
20use serde_json::Value;
21use serde_json::json;
22use simulacrum::AdvanceEpochConfig;
23use simulacrum::Simulacrum;
24use sui_futures::service::Service;
25use sui_indexer_alt::BootstrapGenesis;
26use sui_indexer_alt::config::IndexerConfig;
27use sui_indexer_alt::setup_indexer;
28use sui_indexer_alt_consistent_api::proto::rpc::consistent::v1alpha::AvailableRangeRequest;
29use sui_indexer_alt_consistent_api::proto::rpc::consistent::v1alpha::consistent_service_client::ConsistentServiceClient;
30use sui_indexer_alt_consistent_store::args::RpcArgs as ConsistentArgs;
31use sui_indexer_alt_consistent_store::args::TlsArgs as ConsistentTlsArgs;
32use sui_indexer_alt_consistent_store::config::ServiceConfig as ConsistentConfig;
33use sui_indexer_alt_consistent_store::start_service as start_consistent_store;
34use sui_indexer_alt_framework::IndexerArgs;
35use sui_indexer_alt_framework::ingestion::ClientArgs;
36use sui_indexer_alt_framework::ingestion::ingestion_client::IngestionClientArgs;
37use sui_indexer_alt_framework::pipeline::CommitterConfig;
38use sui_indexer_alt_framework::postgres::schema::watermarks;
39use sui_indexer_alt_graphql::RpcArgs as GraphQlArgs;
40use sui_indexer_alt_graphql::args::SubscriptionArgs;
41use sui_indexer_alt_graphql::config::RpcConfig as GraphQlConfig;
42use sui_indexer_alt_graphql::start_rpc as start_graphql;
43use sui_indexer_alt_jsonrpc::NodeArgs as JsonRpcNodeArgs;
44use sui_indexer_alt_jsonrpc::RpcArgs as JsonRpcArgs;
45use sui_indexer_alt_jsonrpc::config::RpcConfig as JsonRpcConfig;
46use sui_indexer_alt_jsonrpc::start_rpc as start_jsonrpc;
47use sui_indexer_alt_reader::consistent_reader::ConsistentReaderArgs;
48use sui_indexer_alt_reader::fullnode_client::FullnodeArgs;
49use sui_indexer_alt_reader::kv_loader::KvArgs;
50use sui_indexer_alt_reader::system_package_task::SystemPackageTaskArgs;
51use sui_kv_rpc::KvRpcConfig;
52use sui_kv_rpc::KvRpcServer;
53use sui_kvstore::ALL_PIPELINE_NAMES;
54use sui_kvstore::BigTableClient;
55use sui_kvstore::BigTableIndexer;
56use sui_kvstore::CHECKPOINTS_PIPELINE;
57use sui_kvstore::IndexerConfig as BtIndexerConfig;
58use sui_kvstore::IngestionConfig as BtIngestionConfig;
59use sui_kvstore::KeyValueStoreReader;
60use sui_kvstore::PipelineLayer;
61use sui_kvstore::testing::BigTableEmulator;
62use sui_kvstore::testing::INSTANCE_ID;
63use sui_kvstore::testing::create_tables;
64use sui_pg_db::Db;
65use sui_pg_db::DbArgs;
66use sui_pg_db::temp::TempDb;
67use sui_pg_db::temp::get_available_port;
68use sui_protocol_config::Chain;
69use sui_rpc::field::FieldMask;
70use sui_rpc::field::FieldMaskUtil;
71use sui_rpc::merge::Merge;
72use sui_rpc::proto::sui::rpc;
73use sui_types::base_types::ObjectRef;
74use sui_types::base_types::SuiAddress;
75use sui_types::crypto::AccountKeyPair;
76use sui_types::effects::TransactionEffects;
77use sui_types::error::ExecutionError;
78use sui_types::full_checkpoint_content::Checkpoint;
79use sui_types::messages_checkpoint::VerifiedCheckpoint;
80use sui_types::transaction::Transaction;
81use tempfile::TempDir;
82use tokio::time::error::Elapsed;
83use tokio::time::interval;
84use tokio::try_join;
85use url::Url;
86
87pub mod coin_registry;
88pub mod find;
89pub mod graphql;
90pub mod move_helpers;
91pub mod transaction;
92
93/// A simulation of the network, accompanied by off-chain services (database, indexer, RPC),
94/// connected by local data ingestion.
95pub struct FullCluster {
96    /// A simulation of the network, executing transactions and producing checkpoints.
97    executor: Simulacrum,
98
99    /// The off-chain services (database, indexer, RPC) that are ingesting data from the
100    /// simulation.
101    offchain: OffchainCluster,
102
103    /// Temporary directory to store checkpoint information in, so that the indexer can pick it up.
104    #[allow(unused)]
105    temp_dir: TempDir,
106}
107
108/// A collection of the off-chain services (an indexer, a database, and JSON-RPC/GraphQL servers
109/// that read from that database), grouped together to simplify set-up and tear-down for tests. The
110/// included RPC servers do not support transaction dry run and execution.
111///
112/// The database is temporary, and will be cleaned up when the cluster is dropped, and the RPCs are
113/// set-up to listen on a random, available port, to avoid conflicts when multiple instances are
114/// running concurrently in the same process.
115pub struct OffchainCluster {
116    /// The address the consistent store is listening on.
117    consistent_listen_address: SocketAddr,
118
119    /// The address the JSON-RPC server is listening on.
120    jsonrpc_listen_address: SocketAddr,
121
122    /// The address the GraphQL server is listening on.
123    graphql_listen_address: SocketAddr,
124
125    /// The address the kv-rpc (LedgerService) server is listening on.
126    kv_rpc_listen_address: SocketAddr,
127
128    /// The address kv-rpc's second, unencrypted listener is listening on, when
129    /// `OffchainClusterConfig::kv_rpc_plaintext_listener` is set.
130    kv_rpc_plaintext_listen_address: Option<SocketAddr>,
131
132    /// Read access to BigTable.
133    bigtable_client: BigTableClient,
134
135    /// Read access to the temporary database.
136    db: Db,
137
138    /// The pipelines that the indexer is populating.
139    pipelines: Vec<&'static str>,
140
141    /// Handles to all running services. Held on to so the services are not dropped (and therefore
142    /// aborted) until the cluster is stopped.
143    #[allow(unused)]
144    services: Service,
145
146    /// Handle to the BigTable emulator process.
147    #[allow(unused)]
148    bigtable_emulator: BigTableEmulator,
149
150    /// Hold on to the database so it doesn't get dropped until the cluster is stopped.
151    #[allow(unused)]
152    database: TempDb,
153
154    /// Hold on to the temporary directory where the consistent store writes its data, so it
155    /// doesn't get cleaned up until the cluster is stopped.
156    #[allow(unused)]
157    dir: TempDir,
158}
159
160pub struct OffchainClusterConfig {
161    pub indexer_args: IndexerArgs,
162    pub consistent_indexer_args: IndexerArgs,
163    pub fullnode_args: FullnodeArgs,
164    pub indexer_config: IndexerConfig,
165    pub consistent_config: ConsistentConfig,
166    pub jsonrpc_config: JsonRpcConfig,
167    pub jsonrpc_node_args: JsonRpcNodeArgs,
168    pub graphql_config: GraphQlConfig,
169    pub bootstrap_genesis: Option<BootstrapGenesis>,
170    pub kv_rpc_config: KvRpcConfig,
171    /// Per-pipeline overrides (e.g. rate limits) for the BigTable archival indexer.
172    pub bt_pipeline_layer: PipelineLayer,
173    /// When set, kv-rpc also binds a second, unencrypted listener
174    /// (`sui_kv_rpc::ServerConfig::plaintext_address`), reachable via
175    /// `kv_rpc_plaintext_url`, for tests that exercise it directly.
176    pub kv_rpc_plaintext_listener: bool,
177}
178
179impl FullCluster {
180    /// Creates a cluster with a fresh executor where the off-chain services are set up with a
181    /// default configuration.
182    pub async fn new() -> anyhow::Result<Self> {
183        Self::new_with_configs(
184            Simulacrum::new(),
185            OffchainClusterConfig::default(),
186            &prometheus::Registry::new(),
187        )
188        .await
189    }
190
191    /// Creates a new cluster executing transactions using `executor`. The indexer is configured
192    /// using `indexer_args` and `indexer_config, the JSON-RPC server is configured using
193    /// `jsonrpc_config`, and the GraphQL server is configured using `graphql_config`.
194    pub async fn new_with_configs(
195        mut executor: Simulacrum,
196        offchain_cluster_config: OffchainClusterConfig,
197        registry: &prometheus::Registry,
198    ) -> anyhow::Result<Self> {
199        let (client_args, temp_dir) = local_ingestion_client_args();
200        executor.set_data_ingestion_path(temp_dir.path().to_owned());
201
202        let offchain = OffchainCluster::new(client_args, offchain_cluster_config, registry)
203            .await
204            .context("Failed to create off-chain cluster")?;
205
206        Ok(Self {
207            executor,
208            offchain,
209            temp_dir,
210        })
211    }
212
213    /// Return the reference gas price for the current epoch
214    pub fn reference_gas_price(&self) -> u64 {
215        self.executor.reference_gas_price()
216    }
217
218    /// Create a new account and credit it with `amount` gas units from a faucet account. Returns
219    /// the account, its keypair, and a reference to the gas object it was funded with.
220    pub fn funded_account(
221        &mut self,
222        amount: u64,
223    ) -> anyhow::Result<(SuiAddress, AccountKeyPair, ObjectRef)> {
224        self.executor.funded_account(amount)
225    }
226
227    /// Request gas from the faucet, sent to `address`. Return the object reference of the gas
228    /// object that was sent.
229    pub fn request_gas(
230        &mut self,
231        address: SuiAddress,
232        amount: u64,
233    ) -> anyhow::Result<TransactionEffects> {
234        self.executor.request_gas(address, amount)
235    }
236
237    /// Execute a signed transaction, returning its effects.
238    pub fn execute_transaction(
239        &mut self,
240        tx: Transaction,
241    ) -> anyhow::Result<(TransactionEffects, Option<ExecutionError>)> {
242        self.executor.execute_transaction(tx)
243    }
244
245    /// Execute a system transaction advancing the lock by the given `duration`.
246    pub fn advance_clock(&mut self, duration: Duration) -> TransactionEffects {
247        self.executor.advance_clock(duration)
248    }
249
250    /// Advance the executor into the next epoch. This executes an end-of-epoch transaction and
251    /// creates the epoch's final checkpoint, but does not wait for the off-chain services to ingest
252    /// it — follow with [`create_checkpoint`](Self::create_checkpoint) to sync.
253    pub fn advance_epoch(&mut self) {
254        self.executor.advance_epoch(AdvanceEpochConfig::default());
255    }
256
257    /// Create a new checkpoint containing the transactions executed since the last checkpoint that
258    /// was created, and wait for the off-chain services to ingest it. Returns the checkpoint
259    /// contents.
260    pub async fn create_checkpoint(&mut self) -> VerifiedCheckpoint {
261        let checkpoint = self.executor.create_checkpoint();
262        let timeout = Duration::from_secs(100);
263        let indexer = self
264            .offchain
265            .wait_for_indexer(checkpoint.sequence_number, timeout);
266        let consistent_store = self
267            .offchain
268            .wait_for_consistent_store(checkpoint.sequence_number, timeout);
269        let graphql = self
270            .offchain
271            .wait_for_graphql(checkpoint.sequence_number, timeout);
272        let bigtable = self.offchain.wait_for_bigtable(
273            &ALL_PIPELINE_NAMES,
274            checkpoint.sequence_number,
275            timeout,
276        );
277
278        try_join!(indexer, consistent_store, graphql, bigtable)
279            .expect("Timed out waiting for off-chain services");
280
281        checkpoint
282    }
283
284    /// Unlike [`create_checkpoint`](Self::create_checkpoint), only waits for the base
285    /// `checkpoints` BigTable pipeline to catch up — not the indexer, consistent store, GraphQL,
286    /// or the list-index BigTable pipelines (`tx_seq_digest`, `transaction_bitmap_index`,
287    /// `event_bitmap_index`). Used by tests that need to observe a checkpoint the base pipeline
288    /// has indexed before the (typically throttled, via `bt_pipeline_layer`) list-index
289    /// pipelines have processed it, without unrelated services' sync time giving the throttled
290    /// pipelines room to catch up anyway.
291    pub async fn create_checkpoint_before_list_apis_sync(&mut self) -> VerifiedCheckpoint {
292        let checkpoint = self.executor.create_checkpoint();
293        let timeout = Duration::from_secs(100);
294        self.offchain
295            .wait_for_bigtable(&[CHECKPOINTS_PIPELINE], checkpoint.sequence_number, timeout)
296            .await
297            .expect("Timed out waiting for the base checkpoints pipeline");
298
299        checkpoint
300    }
301
302    /// Waits until every pipeline in `pipelines` has caught up to the given `checkpoint`, or the
303    /// `timeout` is reached (an error).
304    pub async fn wait_for_bigtable(
305        &self,
306        pipelines: &[&str],
307        checkpoint: u64,
308        timeout: Duration,
309    ) -> Result<(), Elapsed> {
310        self.offchain
311            .wait_for_bigtable(pipelines, checkpoint, timeout)
312            .await
313    }
314
315    /// The URL to talk to the database on.
316    pub fn db_url(&self) -> Url {
317        self.offchain.db_url()
318    }
319
320    /// The URL to send Consistent Store requests to.
321    pub fn consistent_store_url(&self) -> Url {
322        self.offchain.consistent_store_url()
323    }
324
325    /// The URL to send JSON-RPC requests to.
326    pub fn jsonrpc_url(&self) -> Url {
327        self.offchain.jsonrpc_url()
328    }
329
330    /// The URL to send GraphQL requests to.
331    pub fn graphql_url(&self) -> Url {
332        self.offchain.graphql_url()
333    }
334
335    /// The URL to send kv-rpc (LedgerService) requests to.
336    pub fn kv_rpc_url(&self) -> Url {
337        self.offchain.kv_rpc_url()
338    }
339
340    /// The URL to send requests to kv-rpc's second, unencrypted listener, when
341    /// `OffchainClusterConfig::kv_rpc_plaintext_listener` was set.
342    pub fn kv_rpc_plaintext_url(&self) -> Option<Url> {
343        self.offchain.kv_rpc_plaintext_url()
344    }
345
346    /// Returns the latest checkpoint that we have all data for in the database, according to the
347    /// watermarks table. Returns `None` if any of the expected pipelines are missing data.
348    pub async fn latest_checkpoint(&self) -> anyhow::Result<Option<u64>> {
349        self.offchain.latest_checkpoint().await
350    }
351
352    /// Waits until the indexer has caught up to the given `checkpoint`, or the `timeout` is
353    /// reached (an error).
354    pub async fn wait_for_indexer(
355        &self,
356        checkpoint: u64,
357        timeout: Duration,
358    ) -> Result<(), Elapsed> {
359        self.offchain.wait_for_indexer(checkpoint, timeout).await
360    }
361
362    /// Waits until the indexer's pruner has caught up to the given `checkpoint`, for the given
363    /// `pipeline`, or the `timeout` is reached (an error).
364    pub async fn wait_for_pruner(
365        &self,
366        pipeline: &str,
367        checkpoint: u64,
368        timeout: Duration,
369    ) -> Result<(), Elapsed> {
370        self.offchain
371            .wait_for_pruner(pipeline, checkpoint, timeout)
372            .await
373    }
374
375    /// Waits until GraphQL has caught up to the given `checkpoint`, or the `timeout` is
376    /// reached (an error).
377    pub async fn wait_for_graphql(
378        &self,
379        checkpoint: u64,
380        timeout: Duration,
381    ) -> Result<(), Elapsed> {
382        self.offchain.wait_for_graphql(checkpoint, timeout).await
383    }
384}
385
386impl OffchainCluster {
387    /// Construct a new off-chain cluster and spin up its constituent services.
388    ///
389    /// - `indexer_args`, `client_args`, and `indexer_config` control the indexer. In particular
390    ///   `client_args` is used to configure the client that the indexer uses to fetch checkpoints.
391    /// - `jsonrpc_config` controls the JSON-RPC server.
392    /// - `graphql_config` controls the GraphQL server.
393    /// - `registry` is used to register metrics for the indexer, JSON-RPC, and GraphQL servers.
394    pub async fn new(
395        client_args: ClientArgs,
396        OffchainClusterConfig {
397            indexer_args,
398            consistent_indexer_args,
399            fullnode_args,
400            indexer_config,
401            consistent_config,
402            jsonrpc_config,
403            jsonrpc_node_args,
404            graphql_config,
405            bootstrap_genesis,
406            kv_rpc_config,
407            bt_pipeline_layer,
408            kv_rpc_plaintext_listener,
409        }: OffchainClusterConfig,
410        registry: &prometheus::Registry,
411    ) -> anyhow::Result<Self> {
412        let consistent_port = get_available_port();
413        let consistent_listen_address =
414            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), consistent_port);
415
416        let jsonrpc_port = get_available_port();
417        let jsonrpc_listen_address = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), jsonrpc_port);
418
419        let graphql_port = get_available_port();
420        let graphql_listen_address = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), graphql_port);
421
422        let kv_rpc_port = get_available_port();
423        let kv_rpc_address = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), kv_rpc_port);
424
425        let kv_rpc_plaintext_address = kv_rpc_plaintext_listener.then(|| {
426            let port = get_available_port();
427            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port)
428        });
429
430        let database = TempDb::new().context("Failed to create database")?;
431        let database_url = database.database().url();
432
433        let dir = tempfile::tempdir().context("Failed to create temporary directory")?;
434        let rocksdb_path = dir.path().join("rocksdb");
435
436        let consistent_args = ConsistentArgs {
437            rpc_listen_address: consistent_listen_address,
438            tls: ConsistentTlsArgs::default(),
439        };
440
441        let jsonrpc_args = JsonRpcArgs {
442            rpc_listen_address: jsonrpc_listen_address,
443            ..Default::default()
444        };
445
446        let graphql_args = GraphQlArgs {
447            rpc_listen_address: graphql_listen_address,
448            no_ide: true,
449        };
450
451        let db = Db::for_read(database_url.clone(), DbArgs::default())
452            .await
453            .context("Failed to connect to database")?;
454
455        let indexer = setup_indexer(
456            database_url.clone(),
457            DbArgs::default(),
458            indexer_args,
459            client_args.clone(),
460            indexer_config,
461            bootstrap_genesis,
462            registry,
463        )
464        .await
465        .context("Failed to setup indexer")?;
466
467        let pipelines: Vec<_> = indexer.pipelines().collect();
468        let indexer = indexer.run().await.context("Failed to start indexer")?;
469
470        let consistent_store = start_consistent_store(
471            rocksdb_path,
472            consistent_indexer_args,
473            client_args.clone(),
474            consistent_args,
475            "0.0.0",
476            consistent_config,
477            registry,
478        )
479        .await
480        .context("Failed to start Consistent Store")?;
481
482        let consistent_reader_args = ConsistentReaderArgs {
483            consistent_store_url: Some(
484                Url::parse(&format!("http://{consistent_listen_address}")).unwrap(),
485            ),
486            ..Default::default()
487        };
488
489        // One switch drives both sides: the kv-rpc server only serves the List
490        // APIs when they are enabled, and the graphql/jsonrpc readers only
491        // consume them when they are. Off by default, matching production.
492        let enable_list_apis = kv_rpc_config.enable_list_apis();
493
494        let (bigtable_client, bigtable_emulator, archival_service) = start_archival(
495            client_args.clone(),
496            kv_rpc_address,
497            kv_rpc_plaintext_address,
498            kv_rpc_config,
499            bt_pipeline_layer,
500            registry,
501        )
502        .await?;
503
504        let kv_args = KvArgs {
505            ledger_grpc_url: Some(
506                format!("http://{kv_rpc_address}")
507                    .parse()
508                    .expect("Failed to parse kv-rpc URI"),
509            ),
510            enable_list_apis: Some(enable_list_apis),
511            ..Default::default()
512        };
513
514        let jsonrpc = start_jsonrpc(
515            Some(database_url.clone()),
516            DbArgs::default(),
517            kv_args.clone(),
518            consistent_reader_args.clone(),
519            jsonrpc_args,
520            jsonrpc_node_args,
521            SystemPackageTaskArgs::default(),
522            jsonrpc_config,
523            registry,
524        )
525        .await
526        .context("Failed to start JSON-RPC server")?;
527
528        let graphql = start_graphql(
529            Some(database_url.clone()),
530            fullnode_args,
531            DbArgs::default(),
532            kv_args,
533            consistent_reader_args,
534            graphql_args,
535            SystemPackageTaskArgs::default(),
536            SubscriptionArgs::default(),
537            "0.0.0",
538            graphql_config,
539            pipelines.iter().map(|p| p.to_string()).collect(),
540            registry,
541        )
542        .await
543        .context("Failed to start GraphQL server")?;
544
545        let services = indexer
546            .merge(consistent_store)
547            .merge(jsonrpc)
548            .merge(graphql)
549            .merge(archival_service);
550
551        Ok(Self {
552            consistent_listen_address,
553            jsonrpc_listen_address,
554            graphql_listen_address,
555            kv_rpc_listen_address: kv_rpc_address,
556            kv_rpc_plaintext_listen_address: kv_rpc_plaintext_address,
557            bigtable_client,
558            db,
559            pipelines,
560            services,
561            bigtable_emulator,
562            database,
563            dir,
564        })
565    }
566
567    /// The URL to talk to the database on.
568    pub fn db_url(&self) -> Url {
569        self.database.database().url().clone()
570    }
571
572    /// The URL to send Consistent Store requests to.
573    pub fn consistent_store_url(&self) -> Url {
574        Url::parse(&format!("http://{}/", self.consistent_listen_address))
575            .expect("Failed to parse RPC URL")
576    }
577
578    /// The URL to send JSON-RPC requests to.
579    pub fn jsonrpc_url(&self) -> Url {
580        Url::parse(&format!("http://{}/", self.jsonrpc_listen_address))
581            .expect("Failed to parse RPC URL")
582    }
583
584    /// The URL to send GraphQL requests to.
585    pub fn graphql_url(&self) -> Url {
586        Url::parse(&format!("http://{}/graphql", self.graphql_listen_address))
587            .expect("Failed to parse RPC URL")
588    }
589
590    /// The URL to send kv-rpc (LedgerService) requests to.
591    pub fn kv_rpc_url(&self) -> Url {
592        Url::parse(&format!("http://{}/", self.kv_rpc_listen_address))
593            .expect("Failed to parse RPC URL")
594    }
595
596    /// The URL to send requests to kv-rpc's second, unencrypted listener, when
597    /// `OffchainClusterConfig::kv_rpc_plaintext_listener` was set.
598    pub fn kv_rpc_plaintext_url(&self) -> Option<Url> {
599        let address = self.kv_rpc_plaintext_listen_address?;
600        Some(Url::parse(&format!("http://{address}/")).expect("Failed to parse RPC URL"))
601    }
602
603    /// Returns the latest checkpoint that we have all data for in the database, according to the
604    /// watermarks table. Returns `None` if any of the expected pipelines are missing data.
605    pub async fn latest_checkpoint(&self) -> anyhow::Result<Option<u64>> {
606        use watermarks::dsl as w;
607
608        let mut conn = self
609            .db
610            .connect()
611            .await
612            .context("Failed to connect to database")?;
613
614        let latest: HashMap<String, i64> = w::watermarks
615            .select((w::pipeline, w::checkpoint_hi_inclusive))
616            .filter(w::pipeline.eq_any(&self.pipelines))
617            .filter(w::reader_lo.le(w::checkpoint_hi_inclusive))
618            .load(&mut conn)
619            .await?
620            .into_iter()
621            .collect();
622
623        if latest.len() != self.pipelines.len() {
624            return Ok(None);
625        }
626
627        Ok(latest.into_values().min().map(|l| l as u64))
628    }
629
630    /// Returns the latest checkpoint that the pruner is willing to prune up to for the given
631    /// `pipeline`.
632    pub async fn latest_pruner_checkpoint(&self, pipeline: &str) -> anyhow::Result<Option<u64>> {
633        use watermarks::dsl as w;
634
635        let mut conn = self
636            .db
637            .connect()
638            .await
639            .context("Failed to connect to database")?;
640
641        let latest: Option<i64> = w::watermarks
642            .select(w::reader_lo)
643            .filter(w::pipeline.eq(pipeline))
644            .first(&mut conn)
645            .await
646            .optional()?;
647
648        Ok(latest.map(|l| l as u64))
649    }
650
651    /// Returns the latest checkpoint that the consistent store is aware of.
652    pub async fn latest_consistent_store_checkpoint(&self) -> anyhow::Result<u64> {
653        ConsistentServiceClient::connect(self.consistent_store_url().to_string())
654            .await
655            .context("Failed to connect to Consistent Store")?
656            .available_range(AvailableRangeRequest {})
657            .await
658            .context("Failed to fetch available range from Consistent Store")?
659            .into_inner()
660            .max_checkpoint
661            .context("Consistent Store has not started yet")
662    }
663
664    /// Returns the latest checkpoint that the GraphQL service is aware of.
665    pub async fn latest_graphql_checkpoint(&self) -> anyhow::Result<u64> {
666        let query = json!({
667            "query": "query { checkpoint { sequenceNumber } }"
668        });
669
670        let client = Client::new();
671        let request = client.post(self.graphql_url()).json(&query);
672        let response = request
673            .send()
674            .await
675            .context("Request to GraphQL server failed")?;
676
677        let body: Value = response
678            .json()
679            .await
680            .context("Failed to parse GraphQL response")?;
681
682        let sequence_number = body
683            .pointer("/data/checkpoint/sequenceNumber")
684            .context("Failed to find checkpoint sequence number in response")?;
685
686        let sequence_number: i64 = serde_json::from_value(sequence_number.clone())
687            .context("Failed to parse sequence number as i64")?;
688
689        ensure!(sequence_number != i64::MAX, "Indexer has not started yet");
690
691        Ok(sequence_number as u64)
692    }
693
694    /// Returns the latest epoch that the GraphQL service is aware of.
695    pub async fn latest_graphql_epoch(&self) -> anyhow::Result<u64> {
696        let query = json!({
697            "query": "query { epoch { epochId } }"
698        });
699
700        let client = Client::new();
701        let request = client.post(self.graphql_url()).json(&query);
702        let response = request
703            .send()
704            .await
705            .context("Request to GraphQL server failed")?;
706
707        let body: Value = response
708            .json()
709            .await
710            .context("Failed to parse GraphQL response")?;
711
712        let epoch_id = body
713            .pointer("/data/epoch/epochId")
714            .context("Failed to find epochId in response")?;
715
716        let epoch_id: i64 =
717            serde_json::from_value(epoch_id.clone()).context("Failed to parse epochId as i64")?;
718
719        ensure!(epoch_id != i64::MAX, "Indexer has not started yet");
720
721        Ok(epoch_id as u64)
722    }
723
724    /// Waits until the indexer has caught up to the given `checkpoint`, or the `timeout` is
725    /// reached (an error).
726    pub async fn wait_for_indexer(
727        &self,
728        checkpoint: u64,
729        timeout: Duration,
730    ) -> Result<(), Elapsed> {
731        tokio::time::timeout(timeout, async move {
732            let mut interval = interval(Duration::from_millis(200));
733            loop {
734                interval.tick().await;
735                if matches!(self.latest_checkpoint().await, Ok(Some(l)) if l >= checkpoint) {
736                    break;
737                }
738            }
739        })
740        .await
741    }
742
743    /// Waits until the indexer's pruner has caught up to the given `checkpoint`, for the given
744    /// `pipeline`, or the `timeout` is reached (an error).
745    pub async fn wait_for_pruner(
746        &self,
747        pipeline: &str,
748        checkpoint: u64,
749        timeout: Duration,
750    ) -> Result<(), Elapsed> {
751        tokio::time::timeout(timeout, async move {
752            let mut interval = interval(Duration::from_millis(200));
753            loop {
754                interval.tick().await;
755                if matches!(self.latest_pruner_checkpoint(pipeline).await, Ok(Some(l)) if l >= checkpoint) {
756                    break;
757                }
758            }
759        }).await
760    }
761
762    /// Waits until the Consistent Store has caught up to the given `checkpoint`, or the `timeout`
763    /// is reached (an error).
764    pub async fn wait_for_consistent_store(
765        &self,
766        checkpoint: u64,
767        timeout: Duration,
768    ) -> Result<(), Elapsed> {
769        tokio::time::timeout(timeout, async move {
770            let mut interval = interval(Duration::from_millis(200));
771            loop {
772                interval.tick().await;
773                if matches!(self.latest_consistent_store_checkpoint().await, Ok(l) if l >= checkpoint) {
774                    break;
775                }
776            }
777        })
778        .await
779    }
780
781    /// Waits until GraphQL has caught up to the given `checkpoint`, or the `timeout` is reached
782    /// (an error).
783    pub async fn wait_for_graphql(
784        &self,
785        checkpoint: u64,
786        timeout: Duration,
787    ) -> Result<(), Elapsed> {
788        tokio::time::timeout(timeout, async move {
789            let mut interval = interval(Duration::from_millis(200));
790            loop {
791                interval.tick().await;
792                if matches!(self.latest_graphql_checkpoint().await, Ok(l) if l >= checkpoint) {
793                    break;
794                }
795            }
796        })
797        .await
798    }
799
800    /// Waits until every pipeline in `pipelines` has caught up to the given `checkpoint`, or the
801    /// `timeout` is reached (an error).
802    pub async fn wait_for_bigtable(
803        &self,
804        pipelines: &[&str],
805        checkpoint: u64,
806        timeout: Duration,
807    ) -> Result<(), Elapsed> {
808        let mut client = self.bigtable_client.clone();
809        tokio::time::timeout(timeout, async move {
810            let mut interval = interval(Duration::from_millis(200));
811            loop {
812                interval.tick().await;
813                if client
814                    .get_watermark_for_pipelines(pipelines)
815                    .await
816                    .is_ok_and(|wm| {
817                        wm.is_some_and(|wm| {
818                            wm.checkpoint_hi_inclusive
819                                .is_some_and(|cp| cp >= checkpoint)
820                        })
821                    })
822                {
823                    break;
824                }
825            }
826        })
827        .await
828    }
829}
830
831impl Default for OffchainClusterConfig {
832    fn default() -> Self {
833        Self {
834            indexer_args: Default::default(),
835            consistent_indexer_args: Default::default(),
836            fullnode_args: FullnodeArgs::default(),
837            indexer_config: IndexerConfig::for_test(),
838            consistent_config: ConsistentConfig::for_test(),
839            jsonrpc_config: Default::default(),
840            jsonrpc_node_args: Default::default(),
841            graphql_config: Default::default(),
842            bootstrap_genesis: None,
843            kv_rpc_config: KvRpcConfig::default(),
844            bt_pipeline_layer: PipelineLayer::default(),
845            kv_rpc_plaintext_listener: false,
846        }
847    }
848}
849
850/// Returns ClientArgs that use a temporary local ingestion path and the TempDir of that path.
851pub fn local_ingestion_client_args() -> (ClientArgs, TempDir) {
852    let temp_dir = tempfile::tempdir()
853        .context("Failed to create data ingestion path")
854        .unwrap();
855    let client_args = ClientArgs {
856        ingestion: IngestionClientArgs {
857            local_ingestion_path: Some(temp_dir.path().to_owned()),
858            ..Default::default()
859        },
860        ..Default::default()
861    };
862    (client_args, temp_dir)
863}
864
865/// Writes a checkpoint file to the given path.
866pub async fn write_checkpoint(path: &Path, checkpoint: Checkpoint) -> anyhow::Result<()> {
867    let sequence_number = checkpoint.summary.sequence_number;
868
869    let mask = FieldMask::from_paths([
870        rpc::v2::Checkpoint::path_builder().sequence_number(),
871        rpc::v2::Checkpoint::path_builder().summary().bcs().value(),
872        rpc::v2::Checkpoint::path_builder().signature().finish(),
873        rpc::v2::Checkpoint::path_builder().contents().bcs().value(),
874        rpc::v2::Checkpoint::path_builder()
875            .transactions()
876            .transaction()
877            .bcs()
878            .value(),
879        rpc::v2::Checkpoint::path_builder()
880            .transactions()
881            .effects()
882            .bcs()
883            .value(),
884        rpc::v2::Checkpoint::path_builder()
885            .transactions()
886            .effects()
887            .unchanged_loaded_runtime_objects()
888            .finish(),
889        rpc::v2::Checkpoint::path_builder()
890            .transactions()
891            .events()
892            .bcs()
893            .value(),
894        rpc::v2::Checkpoint::path_builder()
895            .objects()
896            .objects()
897            .bcs()
898            .value(),
899    ]);
900
901    let proto_checkpoint = rpc::v2::Checkpoint::merge_from(&checkpoint, &mask.into());
902    let proto_bytes = proto_checkpoint.encode_to_vec();
903    let compressed = zstd::encode_all(&proto_bytes[..], 3)?;
904
905    let file_name = format!("{}.binpb.zst", sequence_number);
906    let file_path = path.join(file_name);
907    fs::write(file_path, compressed)?;
908    Ok(())
909}
910
911/// Start the archival stack: BigTable emulator, BigTable indexer, and sui-kv-rpc.
912async fn start_archival(
913    client_args: ClientArgs,
914    kv_rpc_address: SocketAddr,
915    kv_rpc_plaintext_address: Option<SocketAddr>,
916    kv_rpc_config: KvRpcConfig,
917    bt_pipeline_layer: PipelineLayer,
918    registry: &prometheus::Registry,
919) -> anyhow::Result<(BigTableClient, BigTableEmulator, Service)> {
920    let emulator = tokio::task::spawn_blocking(BigTableEmulator::start)
921        .await
922        .context("spawn_blocking panicked")?
923        .context("Failed to start BigTable emulator")?;
924
925    create_tables(emulator.host(), INSTANCE_ID)
926        .await
927        .context("Failed to create BigTable tables")?;
928
929    let bigtable_client =
930        BigTableClient::new_local(emulator.host().to_string(), INSTANCE_ID.to_string())
931            .await
932            .context("Failed to create BigTable client")?;
933
934    let indexer_client =
935        BigTableClient::new_local(emulator.host().to_string(), INSTANCE_ID.to_string())
936            .await
937            .context("Failed to create BigTable client for indexer")?;
938    let bt_indexer = BigTableIndexer::new(
939        indexer_client,
940        IndexerArgs::default(),
941        client_args,
942        BtIngestionConfig::default(),
943        CommitterConfig::default(),
944        BtIndexerConfig::default(),
945        bt_pipeline_layer,
946        Chain::Unknown,
947        registry,
948    )
949    .await
950    .context("Failed to create BigTable indexer")?;
951
952    // Use the BigTable wrapper, not the raw framework indexer, so bitmap
953    // committer background tasks are supervised for the duration of the test.
954    let bt_indexer_service = bt_indexer
955        .run()
956        .await
957        .context("Failed to start BigTable indexer")?;
958
959    let kv_rpc_server = KvRpcServer::new_local_with_config(
960        emulator.host().to_string(),
961        INSTANCE_ID.to_string(),
962        None,
963        kv_rpc_config.ledger_history(),
964        kv_rpc_config.request_bigtable_concurrency(),
965        kv_rpc_config.stages(),
966        kv_rpc_config.enable_list_apis(),
967    )
968    .await
969    .context("Failed to create KvRpcServer")?;
970    let kv_rpc_service = kv_rpc_server
971        .start_service(
972            kv_rpc_address,
973            sui_kv_rpc::ServerConfig {
974                plaintext_address: kv_rpc_plaintext_address,
975                ..Default::default()
976            },
977        )
978        .await
979        .context("Failed to start kv-rpc server")?;
980
981    let service = bt_indexer_service.merge(kv_rpc_service);
982    Ok((bigtable_client, emulator, service))
983}