From 19373d806d7bedbb8f701d169706d23a3d1c2d8f Mon Sep 17 00:00:00 2001
From: Lin Yihai <1161813899@qq.com>
Date: Fri, 6 Dec 2024 23:02:15 +0800
Subject: [PATCH] chore: Add timeout setting for `find_ttl`. (#5088)

---
 src/mito2/src/compaction.rs | 19 ++++++++++++-------
 src/mito2/src/config.rs     |  3 +++
 src/mito2/src/error.rs      | 10 ++++++++++
 3 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs
index 31e1b0674f72..a4094af74121 100644
--- a/src/mito2/src/compaction.rs
+++ b/src/mito2/src/compaction.rs
@@ -51,6 +51,7 @@ use crate::config::MitoConfig;
 use crate::error::{
     CompactRegionSnafu, Error, GetSchemaMetadataSnafu, RegionClosedSnafu, RegionDroppedSnafu,
     RegionTruncatedSnafu, RemoteCompactionSnafu, Result, TimeRangePredicateOverflowSnafu,
+    TimeoutSnafu,
 };
 use crate::metrics::COMPACTION_STAGE_ELAPSED;
 use crate::read::projection::ProjectionMapper;
@@ -445,13 +446,17 @@ async fn find_ttl(
         return Ok(table_ttl);
     }
 
-    let ttl = schema_metadata_manager
-        .get_schema_options_by_table_id(table_id)
-        .await
-        .context(GetSchemaMetadataSnafu)?
-        .and_then(|options| options.ttl)
-        .unwrap_or_default()
-        .into();
+    let ttl = tokio::time::timeout(
+        crate::config::FETCH_OPTION_TIMEOUT,
+        schema_metadata_manager.get_schema_options_by_table_id(table_id),
+    )
+    .await
+    .context(TimeoutSnafu)?
+    .context(GetSchemaMetadataSnafu)?
+    .and_then(|options| options.ttl)
+    .unwrap_or_default()
+    .into();
+
     Ok(ttl)
 }
 
diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs
index 8cd2b08f2e59..797c42f8084c 100644
--- a/src/mito2/src/config.rs
+++ b/src/mito2/src/config.rs
@@ -45,6 +45,9 @@ const PAGE_CACHE_SIZE_FACTOR: u64 = 8;
 /// Use `1/INDEX_CREATE_MEM_THRESHOLD_FACTOR` of OS memory size as mem threshold for creating index
 const INDEX_CREATE_MEM_THRESHOLD_FACTOR: u64 = 16;
 
+/// Fetch option timeout
+pub(crate) const FETCH_OPTION_TIMEOUT: Duration = Duration::from_secs(10);
+
 /// Configuration for [MitoEngine](crate::engine::MitoEngine).
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
 #[serde(default)]
diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs
index 6cb4f8abdd7a..407c8c29e258 100644
--- a/src/mito2/src/error.rs
+++ b/src/mito2/src/error.rs
@@ -30,6 +30,7 @@ use snafu::{Location, Snafu};
 use store_api::logstore::provider::Provider;
 use store_api::manifest::ManifestVersion;
 use store_api::storage::RegionId;
+use tokio::time::error::Elapsed;
 
 use crate::cache::file_cache::FileType;
 use crate::region::{RegionLeaderState, RegionRoleState};
@@ -877,6 +878,14 @@ pub enum Error {
         #[snafu(implicit)]
         location: Location,
     },
+
+    #[snafu(display("Timeout"))]
+    Timeout {
+        #[snafu(source)]
+        error: Elapsed,
+        #[snafu(implicit)]
+        location: Location,
+    },
 }
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
@@ -1010,6 +1019,7 @@ impl ErrorExt for Error {
             DecodeStats { .. } | StatsNotPresent { .. } => StatusCode::Internal,
             RegionBusy { .. } => StatusCode::RegionBusy,
             GetSchemaMetadata { source, .. } => source.status_code(),
+            Timeout { .. } => StatusCode::Cancelled,
         }
     }