Add new subcommands for forest-cli snasphot

CHANGELOG.md

@@ -29,6 +29,8 @@
 
 ### Added
 
+- [#6082](https://github.com/ChainSafe/forest/issues/6082) Added `forest-cli snapshot export-status` and `forest-cli snapshot export-cancel` subcommands to monitor or cancel an export, respectively.
+
 - [#6166](https://github.com/ChainSafe/forest/pull/6166) Gate `JWT` expiration validation behind environment variable `FOREST_JWT_DISABLE_EXP_VALIDATION`.
 
 ### Changed

scripts/tests/calibnet_export_check.sh

@@ -20,8 +20,6 @@ $FOREST_CLI_PATH snapshot export --tipset 0 --format "$format"
 echo "Exporting zstd compressed snapshot in $format format"
 $FOREST_CLI_PATH snapshot export --format "$format"
 
-$FOREST_CLI_PATH shutdown --force
-
 for f in *.car.zst; do
   echo "Inspecting archive info $f"
   $FOREST_TOOL_PATH archive info "$f"
@@ -39,3 +37,36 @@ for f in *.car.zst; do
   echo "Validating CAR file $f"
   $FOREST_TOOL_PATH snapshot validate "$f"
 done
+
+echo "Cleaning up the last snapshot"
+rm --force --verbose ./*.{car,car.zst,sha256sum}
+
+output=$($FOREST_CLI_PATH snapshot export-status --format json)
+is_exporting=$(echo "$output" | jq -r '.exporting')
+echo "Testing that no export is in progress"
+if [ "$is_exporting" == "true" ]; then
+  exit 1
+fi
+
+echo "Exporting zstd compressed snapshot"
+$FOREST_CLI_PATH snapshot export
+sleep 1
+
+output=$($FOREST_CLI_PATH snapshot export-status --format json)
+is_exporting=$(echo "$output" | jq -r '.exporting')
+echo "Testing that export is in progress"
+if [ "$is_exporting" == "false" ]; then
+  exit 1
+fi
+
+$FOREST_CLI_PATH snapshot export-cancel
+
+output=$($FOREST_CLI_PATH snapshot export-status --format json)
+is_exporting=$(echo "$output" | jq -r '.exporting')
+is_cancelled=$(echo "$output" | jq -r '.cancelled')
+echo "Testing that export has been cancelled"
+if [ "$is_exporting" == "true" ] || [ "$is_cancelled" == "false" ]; then
+  exit 1
+fi
+
+$FOREST_CLI_PATH shutdown --force

src/chain/mod.rs

@@ -161,7 +161,8 @@ async fn export_to_forest_car<D: Digest>(
             tipset.clone().chain_owned(Arc::clone(db)),
             stateroot_lookup_limit,
         )
-        .with_seen(seen),
+        .with_seen(seen)
+        .track_progress(true),
     );
 
     // Encode Ipld key-value pairs in zstd frames

src/cli/subcommands/snapshot_cmd.rs

@@ -7,6 +7,7 @@ use crate::cli_shared::snapshot::{self, TrustedVendor};
 use crate::db::car::forest::new_forest_car_temp_path_in;
 use crate::networks::calibnet;
 use crate::rpc::chain::ForestChainExportDiffParams;
+use crate::rpc::types::ApiExportResult;
 use crate::rpc::{self, chain::ForestChainExportParams, prelude::*};
 use crate::shim::policy::policy_constants::CHAIN_FINALITY;
 use anyhow::Context as _;
@@ -19,6 +20,12 @@ use std::{
 };
 use tokio::io::AsyncWriteExt;
 
+#[derive(Debug, Clone, clap::ValueEnum)]
+pub enum Format {
+    Json,
+    Text,
+}
+
 #[derive(Debug, Subcommand)]
 pub enum SnapshotCommands {
     /// Export a snapshot of the chain to `<output_path>`
@@ -42,6 +49,17 @@ pub enum SnapshotCommands {
         #[arg(long, value_enum, default_value_t = FilecoinSnapshotVersion::V2)]
         format: FilecoinSnapshotVersion,
     },
+    /// Show status of the current export.
+    ExportStatus {
+        /// Wait until it completes and print progress.
+        #[arg(long)]
+        wait: bool,
+        /// Format of the output. `json` or `text`.
+        #[arg(long, value_enum, default_value_t = Format::Text)]
+        format: Format,
+    },
+    /// Cancel the current export.
+    ExportCancel {},
     /// Export a diff snapshot between `from` and `to` epochs to `<output_path>`
     ExportDiff {
         /// `./forest_snapshot_diff_{chain}_{from}_{to}+{depth}.car.zst`.
@@ -152,7 +170,7 @@ impl SnapshotCommands {
 
                 // Manually construct RpcRequest because snapshot export could
                 // take a few hours on mainnet
-                let hash_result = client
+                let export_result = client
                     .call(ForestChainExport::request((params,))?.with_timeout(Duration::MAX))
                     .await?;
 
@@ -161,13 +179,106 @@ impl SnapshotCommands {
                 _ = handle.await;
 
                 if !dry_run {
-                    if let Some(hash) = hash_result {
-                        save_checksum(&output_path, hash).await?;
+                    match export_result.clone() {
+                        ApiExportResult::Done(hash_opt) => {
+                            // Move the file first; prevents orphaned checksum on persist error.
+                            temp_path.persist(&output_path)?;
+                            if let Some(hash) = hash_opt {
+                                save_checksum(&output_path, hash).await?;
+                            }
+                        }
+                        ApiExportResult::Cancelled => { /* no file to persist on cancel */ }
                     }
-                    temp_path.persist(output_path)?;
                 }
 
-                println!("Export completed.");
+                match export_result {
+                    ApiExportResult::Done(_) => {
+                        println!("Export completed.");
+                    }
+                    ApiExportResult::Cancelled => {
+                        println!("Export cancelled.");
+                    }
+                }
+                Ok(())
+            }
+            Self::ExportStatus { wait, format } => {
+                let result = client
+                    .call(
+                        ForestChainExportStatus::request(())?.with_timeout(Duration::from_secs(30)),
+                    )
+                    .await?;
+                if !result.exporting
+                    && let Format::Text = format
+                {
+                    if result.cancelled {
+                        println!("No export in progress (last export was cancelled)");
+                    } else {
+                        println!("No export in progress");
+                    }
+                    return Ok(());
+                }
+                if wait {
+                    let elapsed = chrono::Utc::now()
+                        .signed_duration_since(result.start_time)
+                        .to_std()
+                        .unwrap_or(Duration::ZERO);
+                    let pb = ProgressBar::new(10000)
+                        .with_elapsed(elapsed)
+                        .with_message("Exporting");
+                    pb.set_style(
+                        ProgressStyle::with_template(
+                            "[{elapsed_precise}] [{wide_bar}] {percent}% {msg} ",
+                        )
+                        .expect("indicatif template must be valid")
+                        .progress_chars("#>-"),
+                    );
+                    loop {
+                        let result = client
+                            .call(
+                                ForestChainExportStatus::request(())?
+                                    .with_timeout(Duration::from_secs(30)),
+                            )
+                            .await?;
+                        if result.cancelled {
+                            pb.set_message("Export cancelled");
+                            pb.abandon();
+
+                            return Ok(());
+                        }
+                        let position = (result.progress.clamp(0.0, 1.0) * 10000.0).trunc() as u64;
+                        pb.set_position(position);
+
+                        if position >= 10000 {
+                            break;
+                        }
+                        tokio::time::sleep(Duration::from_millis(500)).await;
+                    }
+                    pb.finish_with_message("Export completed");
+
+                    return Ok(());
+                }
+                match format {
+                    Format::Text => {
+                        println!("Exporting: {:.1}%", result.progress.clamp(0.0, 1.0) * 100.0);
+                    }
+                    Format::Json => {
+                        println!("{}", serde_json::to_string_pretty(&result)?);
+                    }
+                }
+
+                Ok(())
+            }
+            Self::ExportCancel {} => {
+                let result = client
+                    .call(
+                        ForestChainExportCancel::request(())?.with_timeout(Duration::from_secs(30)),
+                    )
+                    .await?;
+                if result {
+                    println!("Export cancelled.");
+                } else {
+                    println!("No export in progress to cancel.");
+                }
                 Ok(())
             }
             Self::ExportDiff {

src/ipld/util.rs

@@ -8,15 +8,58 @@ use crate::shim::clock::ChainEpoch;
 use crate::utils::db::car_stream::CarBlock;
 use crate::utils::encoding::extract_cids;
 use crate::utils::multihash::prelude::*;
+use chrono::{DateTime, Utc};
 use cid::Cid;
 use futures::Stream;
 use fvm_ipld_blockstore::Blockstore;
+use parking_lot::Mutex;
 use pin_project_lite::pin_project;
 use std::borrow::Borrow;
 use std::collections::VecDeque;
 use std::pin::Pin;
+use std::sync::LazyLock;
 use std::task::{Context, Poll};
 
+#[derive(Default)]
+pub struct ExportStatus {
+    pub epoch: i64,
+    pub initial_epoch: i64,
+    pub exporting: bool,
+    pub cancelled: bool,
+    pub start_time: DateTime<Utc>,
+}
+
+pub static CHAIN_EXPORT_STATUS: LazyLock<Mutex<ExportStatus>> =
+    LazyLock::new(|| ExportStatus::default().into());
+
+fn update_epoch(new_value: i64) {
+    let mut mutex = CHAIN_EXPORT_STATUS.lock();
+    mutex.epoch = new_value;
+    if mutex.initial_epoch == 0 {
+        mutex.initial_epoch = new_value;
+    }
+}
+
+pub fn start_export() {
+    let mut mutex = CHAIN_EXPORT_STATUS.lock();
+    mutex.epoch = 0;
+    mutex.initial_epoch = 0;
+    mutex.exporting = true;
+    mutex.cancelled = false;
+    mutex.start_time = Utc::now();
+}
+
+pub fn end_export() {
+    let mut mutex = CHAIN_EXPORT_STATUS.lock();
+    mutex.exporting = false;
+}
+
+pub fn cancel_export() {
+    let mut mutex = CHAIN_EXPORT_STATUS.lock();
+    mutex.exporting = false;
+    mutex.cancelled = true;
+}
+
 fn should_save_block_to_snapshot(cid: Cid) -> bool {
     // Don't include identity CIDs.
     // We only include raw and dagcbor, for now.
@@ -112,6 +155,7 @@ pin_project! {
         seen: CidHashSet,
         stateroot_limit_exclusive: ChainEpoch,
         fail_on_dead_links: bool,
+        track_progress: bool,
     }
 }
 
@@ -126,6 +170,11 @@ impl<DB, T> ChainStream<DB, T> {
         self
     }
 
+    pub fn track_progress(mut self, track_progress: bool) -> Self {
+        self.track_progress = track_progress;
+        self
+    }
+
     #[allow(dead_code)]
     pub fn into_seen(self) -> CidHashSet {
         self.seen
@@ -155,6 +204,7 @@ pub fn stream_chain<DB: Blockstore, T: Borrow<Tipset>, ITER: Iterator<Item = T>
         seen: CidHashSet::default(),
         stateroot_limit_exclusive,
         fail_on_dead_links: true,
+        track_progress: false,
     }
 }
 
@@ -197,6 +247,9 @@ impl<DB: Blockstore, T: Borrow<Tipset>, ITER: Iterator<Item = T> + Unpin> Stream
                         }
                     }
                     Iterate(epoch, block_cid, _type, cid_vec) => {
+                        if *this.track_progress {
+                            update_epoch(*epoch);
+                        }
                         while let Some(cid) = cid_vec.pop_front() {
                             // The link traversal implementation assumes there are three types of encoding:
                             // 1. DAG_CBOR: needs to be reachable, so we add it to the queue and load.
@@ -242,6 +295,9 @@ impl<DB: Blockstore, T: Borrow<Tipset>, ITER: Iterator<Item = T> + Unpin> Stream
                 for block in tipset.borrow().block_headers() {
                     let (cid, data) = block.car_block()?;
                     if this.seen.insert(cid) {
+                        if *this.track_progress {
+                            update_epoch(block.epoch);
+                        }
                         // Make sure we always yield a block otherwise.
                         this.dfs.push_back(Emit(cid, Some(data)));