fix: ensure shard lock is release in corner case

Cargo.toml

@@ -23,7 +23,7 @@ members = [
     "components/arrow_ext",
     "components/bytes_ext",
     "components/codec",
-    "components/future_cancel",
+    "components/future_ext",
     "components/hash_ext",
     "components/id_allocator",
     "components/logger",
@@ -94,7 +94,7 @@ common_types = { path = "common_types" }
 datafusion = { git = "https://github.com/CeresDB/arrow-datafusion.git", rev = "a6dcd943051a083693c352c6b4279156548490a0" }
 datafusion-proto = { git = "https://github.com/CeresDB/arrow-datafusion.git", rev = "a6dcd943051a083693c352c6b4279156548490a0" }
 df_operator = { path = "df_operator" }
-future_cancel = { path = "components/future_cancel" }
+future_ext = { path = "components/future_ext" }
 etcd-client = "0.10.3"
 env_logger = "0.6"
 futures = "0.3"

analytic_engine/Cargo.toml

@@ -26,7 +26,7 @@ ceresdbproto = { workspace = true }
 codec = { workspace = true }
 common_types = { workspace = true }
 datafusion = { workspace = true }
-future_cancel = { workspace = true }
+future_ext = { workspace = true }
 futures = { workspace = true }
 generic_error = { workspace = true }
 hex = { workspace = true }

analytic_engine/src/table/mod.rs

@@ -15,7 +15,7 @@ use common_types::{
     time::TimeRange,
 };
 use datafusion::{common::Column, logical_expr::Expr};
-use future_cancel::CancellationSafeFuture;
+use future_ext::CancellationSafeFuture;
 use futures::TryStreamExt;
 use generic_error::BoxError;
 use log::{error, warn};

cluster/Cargo.toml

@@ -17,6 +17,7 @@ catalog = { workspace = true }
 ceresdbproto = { workspace = true }
 common_types = { workspace = true }
 etcd-client = { workspace = true }
+future_ext = { workspace = true }
 generic_error = { workspace = true }
 log = { workspace = true }
 macros = { workspace = true }

components/future_cancel/Cargo.toml → components/future_ext/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "future_cancel"
+name = "future_ext"
 
 [package.version]
 workspace = true
@@ -13,6 +13,4 @@ workspace = true
 [dependencies]
 futures = { workspace = true }
 runtime = { workspace = true }
-
-[dev-dependencies]
-tokio = { workspace = true }
+tokio = { workspace = true, features = ["time"] }

components/future_ext/src/lib.rs

@@ -0,0 +1,9 @@
+// Copyright 2022-2023 CeresDB Project Authors. Licensed under Apache-2.0.
+
+//! Future extensions.
+
+mod cancel;
+mod retry;
+
+pub use cancel::CancellationSafeFuture;
+pub use retry::{retry_async, RetryConfig};

components/future_ext/src/retry.rs

@@ -0,0 +1,97 @@
+// Copyright 2022-2023 CeresDB Project Authors. Licensed under Apache-2.0.
+
+//! Util function to retry future.
+
+use std::time::Duration;
+
+use futures::Future;
+
+// TODO: add backoff
+// https://github.com/apache/arrow-rs/blob/dfb642809e93c2c1b8343692f4e4b3080000f988/object_store/src/client/backoff.rs#L26
+pub struct RetryConfig {
+    pub max_retries: usize,
+    pub interval: Duration,
+}
+
+impl Default for RetryConfig {
+    fn default() -> Self {
+        Self {
+            max_retries: 3,
+            interval: Duration::from_millis(500),
+        }
+    }
+}
+
+pub async fn retry_async<F, Fut, T, E>(f: F, config: &RetryConfig) -> Fut::Output
+where
+    F: Fn() -> Fut,
+    Fut: Future<Output = Result<T, E>>,
+{
+    for _ in 0..config.max_retries {
+        let result = f().await;
+
+        if result.is_ok() {
+            return result;
+        }
+        tokio::time::sleep(config.interval).await;
+    }
+
+    f().await
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::atomic::{AtomicU8, Ordering};
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_retry_async() {
+        let config = RetryConfig {
+            max_retries: 3,
+            interval: Duration::from_millis(5),
+        };
+
+        // always fails
+        {
+            let runs = AtomicU8::new(0);
+            let f = || {
+                runs.fetch_add(1, Ordering::Relaxed);
+                futures::future::err::<i32, i32>(1)
+            };
+
+            let ret = retry_async(f, &config).await;
+            assert!(ret.is_err());
+            assert_eq!(4, runs.load(Ordering::Relaxed));
+        }
+
+        // succeed directly
+        {
+            let runs = AtomicU8::new(0);
+            let f = || {
+                runs.fetch_add(1, Ordering::Relaxed);
+                futures::future::ok::<i32, i32>(1)
+            };
+
+            let ret = retry_async(f, &config).await;
+            assert_eq!(1, ret.unwrap());
+            assert_eq!(1, runs.load(Ordering::Relaxed));
+        }
+
+        // fail 2 times, then succeed
+        {
+            let runs = AtomicU8::new(0);
+            let f = || {
+                if runs.fetch_add(1, Ordering::Relaxed) < 2 {
+                    return futures::future::err::<_, i32>(1);
+                }
+
+                futures::future::ok::<_, i32>(2)
+            };
+
+            let ret = retry_async(f, &config).await;
+            assert_eq!(2, ret.unwrap());
+            assert_eq!(3, runs.load(Ordering::Relaxed));
+        }
+    }
+}

server/Cargo.toml

@@ -23,6 +23,7 @@ cluster = { workspace = true }
 common_types = { workspace = true }
 datafusion = { workspace = true }
 df_operator = { workspace = true }
+future_ext = { workspace = true }
 futures = { workspace = true }
 generic_error = { workspace = true }
 http = "0.2"

server/src/grpc/meta_event_service/mod.rs

@@ -2,7 +2,10 @@
 
 // Meta event rpc service implementation.
 
-use std::{sync::Arc, time::Instant};
+use std::{
+    sync::Arc,
+    time::{Duration, Instant},
+};
 
 use analytic_engine::setup::OpenedWals;
 use async_trait::async_trait;
@@ -25,6 +28,7 @@ use cluster::{
     ClusterRef,
 };
 use common_types::{schema::SchemaEncoder, table::ShardId};
+use future_ext::RetryConfig;
 use generic_error::BoxError;
 use log::{error, info, warn};
 use meta_client::types::{ShardInfo, TableInfo};
@@ -79,6 +83,12 @@ macro_rules! extract_updated_table_info {
     }};
 }
 
+// TODO: configure retry
+const RETRY: RetryConfig = RetryConfig {
+    max_retries: 10,
+    interval: Duration::from_secs(5),
+};
+
 /// Builder for [MetaServiceImpl].
 pub struct Builder<Q> {
     pub cluster: ClusterRef,
@@ -279,15 +289,24 @@ async fn do_open_shard(ctx: HandlerContext, shard_info: ShardInfo) -> Result<()>
     // Try to lock the shard in node level.
     ctx.acquire_shard_lock(shard_info.id).await?;
 
-    let shard = ctx
-        .cluster
-        .open_shard(&shard_info)
-        .await
-        .box_err()
-        .context(ErrWithCause {
-            code: StatusCode::Internal,
-            msg: "fail to open shards in cluster",
-        })?;
+    // We need to ensure `open_shard` succeeds, otherwise it won't heartbeat to
+    // meta.
+    let shard = future_ext::retry_async(
+        || async {
+            ctx.cluster.open_shard(&shard_info).await.map_err(|e| {
+                error!("Open shard failed, id:{}, err:{e}", shard_info.id);
+                e
+            })
+        },
+        &RETRY,
+    )
+    .await;
+    let shard = match shard {
+        Ok(v) => v,
+        Err(_) => panic!(
+            "Open shard failed and we have to panic otherwise this shard wont be assigned again."
+        ),
+    };
 
     let open_ctx = OpenContext {
         catalog: ctx.default_catalog.clone(),
@@ -297,6 +316,8 @@ async fn do_open_shard(ctx: HandlerContext, shard_info: ShardInfo) -> Result<()>
         engine: ANALYTIC_ENGINE_TYPE.to_string(),
     };
 
+    // This `open` may only open part of tables in this shard, and this is
+    // allowed via shard status(PartialOpen) mechanism.
     shard.open(open_ctx).await.box_err().context(ErrWithCause {
         code: StatusCode::Internal,
         msg: format!("fail to open shard, id:{}", shard_info.id),