1use 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
93pub struct FullCluster {
96 executor: Simulacrum,
98
99 offchain: OffchainCluster,
102
103 #[allow(unused)]
105 temp_dir: TempDir,
106}
107
108pub struct OffchainCluster {
116 consistent_listen_address: SocketAddr,
118
119 jsonrpc_listen_address: SocketAddr,
121
122 graphql_listen_address: SocketAddr,
124
125 kv_rpc_listen_address: SocketAddr,
127
128 kv_rpc_plaintext_listen_address: Option<SocketAddr>,
131
132 bigtable_client: BigTableClient,
134
135 db: Db,
137
138 pipelines: Vec<&'static str>,
140
141 #[allow(unused)]
144 services: Service,
145
146 #[allow(unused)]
148 bigtable_emulator: BigTableEmulator,
149
150 #[allow(unused)]
152 database: TempDb,
153
154 #[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 pub bt_pipeline_layer: PipelineLayer,
173 pub kv_rpc_plaintext_listener: bool,
177}
178
179impl FullCluster {
180 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 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 pub fn reference_gas_price(&self) -> u64 {
215 self.executor.reference_gas_price()
216 }
217
218 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 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 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 pub fn advance_clock(&mut self, duration: Duration) -> TransactionEffects {
247 self.executor.advance_clock(duration)
248 }
249
250 pub fn advance_epoch(&mut self) {
254 self.executor.advance_epoch(AdvanceEpochConfig::default());
255 }
256
257 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 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 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 pub fn db_url(&self) -> Url {
317 self.offchain.db_url()
318 }
319
320 pub fn consistent_store_url(&self) -> Url {
322 self.offchain.consistent_store_url()
323 }
324
325 pub fn jsonrpc_url(&self) -> Url {
327 self.offchain.jsonrpc_url()
328 }
329
330 pub fn graphql_url(&self) -> Url {
332 self.offchain.graphql_url()
333 }
334
335 pub fn kv_rpc_url(&self) -> Url {
337 self.offchain.kv_rpc_url()
338 }
339
340 pub fn kv_rpc_plaintext_url(&self) -> Option<Url> {
343 self.offchain.kv_rpc_plaintext_url()
344 }
345
346 pub async fn latest_checkpoint(&self) -> anyhow::Result<Option<u64>> {
349 self.offchain.latest_checkpoint().await
350 }
351
352 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 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 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 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 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 pub fn db_url(&self) -> Url {
569 self.database.database().url().clone()
570 }
571
572 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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
850pub 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
865pub 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
911async 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 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}