From 96ce9d0b55f32fda46e44d74e5aae51a92fb57f7 Mon Sep 17 00:00:00 2001
From: zwang28 <70626450+zwang28@users.noreply.github.com>
Date: Sun, 25 Aug 2024 16:57:24 +0800
Subject: [PATCH] fix(ctl): support the configuration of read timeout for
 restoring meta backup (#18225)

---
 src/meta/src/backup_restore/restore.rs         | 8 ++++++++
 src/meta/src/backup_restore/restore_impl/v1.rs | 6 +++---
 src/meta/src/backup_restore/restore_impl/v2.rs | 6 +++---
 src/meta/src/backup_restore/utils.rs           | 6 +++++-
 4 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/src/meta/src/backup_restore/restore.rs b/src/meta/src/backup_restore/restore.rs
index 0d0ee5ba3698..00baa5eaf129 100644
--- a/src/meta/src/backup_restore/restore.rs
+++ b/src/meta/src/backup_restore/restore.rs
@@ -79,6 +79,12 @@ pub struct RestoreOpts {
     /// Print the target snapshot, but won't restore to meta store.
     #[clap(long)]
     pub dry_run: bool,
+    /// The read timeout for object store
+    #[clap(long, default_value_t = 600000)]
+    pub read_attempt_timeout_ms: u64,
+    /// The maximum number of read retry attempts for the object store.
+    #[clap(long, default_value_t = 3)]
+    pub read_retry_attempts: u64,
 }
 
 async fn restore_hummock_version(
@@ -251,6 +257,8 @@ mod tests {
             hummock_storage_url: "memory".to_string(),
             hummock_storage_directory: "".to_string(),
             dry_run: false,
+            read_attempt_timeout_ms: 60000,
+            read_retry_attempts: 3,
         }
     }
 
diff --git a/src/meta/src/backup_restore/restore_impl/v1.rs b/src/meta/src/backup_restore/restore_impl/v1.rs
index 303e6e6ad5f5..cc0fdb1baa41 100644
--- a/src/meta/src/backup_restore/restore_impl/v1.rs
+++ b/src/meta/src/backup_restore/restore_impl/v1.rs
@@ -43,7 +43,7 @@ impl Loader<ClusterMetadata> for LoaderV1 {
         let backup_store = &self.backup_store;
         let snapshot_list = &backup_store.manifest().snapshot_metadata;
         let mut target_snapshot: MetaSnapshotV1 = backup_store.get(target_id).await?;
-        tracing::info!(
+        tracing::debug!(
             "snapshot {} before rewrite:\n{}",
             target_id,
             target_snapshot
@@ -78,11 +78,11 @@ impl Loader<ClusterMetadata> for LoaderV1 {
             }
             target_snapshot.metadata.default_cf = newest_snapshot.metadata.default_cf;
             tracing::info!(
-                "snapshot {} after rewrite by snapshot {}:\n{}",
+                "snapshot {} is rewritten by snapshot {}:\n",
                 target_id,
                 newest_id,
-                target_snapshot,
             );
+            tracing::debug!("{target_snapshot}",);
         }
         Ok(target_snapshot)
     }
diff --git a/src/meta/src/backup_restore/restore_impl/v2.rs b/src/meta/src/backup_restore/restore_impl/v2.rs
index a431b455063b..a887293e0c8e 100644
--- a/src/meta/src/backup_restore/restore_impl/v2.rs
+++ b/src/meta/src/backup_restore/restore_impl/v2.rs
@@ -37,7 +37,7 @@ impl Loader<MetadataV2> for LoaderV2 {
     async fn load(&self, target_id: MetaSnapshotId) -> BackupResult<MetaSnapshot<MetadataV2>> {
         let snapshot_list = &self.backup_store.manifest().snapshot_metadata;
         let mut target_snapshot: MetaSnapshotV2 = self.backup_store.get(target_id).await?;
-        tracing::info!(
+        tracing::debug!(
             "snapshot {} before rewrite:\n{}",
             target_id,
             target_snapshot
@@ -73,11 +73,11 @@ impl Loader<MetadataV2> for LoaderV2 {
             }
             target_snapshot.metadata.hummock_sequences = newest_snapshot.metadata.hummock_sequences;
             tracing::info!(
-                "snapshot {} after rewrite by snapshot {}:\n{}",
+                "snapshot {} is rewritten by snapshot {}:\n",
                 target_id,
                 newest_id,
-                target_snapshot,
             );
+            tracing::debug!("{target_snapshot}");
         }
         Ok(target_snapshot)
     }
diff --git a/src/meta/src/backup_restore/utils.rs b/src/meta/src/backup_restore/utils.rs
index 6f6c1dd09ece..4c4c9a336b44 100644
--- a/src/meta/src/backup_restore/utils.rs
+++ b/src/meta/src/backup_restore/utils.rs
@@ -120,11 +120,15 @@ pub async fn get_meta_store(opts: RestoreOpts) -> BackupResult<MetaStoreBackendI
 }
 
 pub async fn get_backup_store(opts: RestoreOpts) -> BackupResult<MetaSnapshotStorageRef> {
+    let mut config = ObjectStoreConfig::default();
+    config.retry.read_attempt_timeout_ms = opts.read_attempt_timeout_ms;
+    config.retry.read_retry_attempts = opts.read_retry_attempts as usize;
+
     let object_store = build_remote_object_store(
         &opts.backup_storage_url,
         Arc::new(ObjectStoreMetrics::unused()),
         "Meta Backup",
-        Arc::new(ObjectStoreConfig::default()),
+        Arc::new(config),
     )
     .await;
     let backup_store =