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 }

cluster/src/cluster_impl.rs

@@ -8,7 +8,6 @@ use std::{
 use async_trait::async_trait;
 use common_types::table::ShardId;
 use etcd_client::ConnectOptions;
-use generic_error::BoxError;
 use log::{error, info, warn};
 use meta_client::{
     types::{
@@ -30,7 +29,7 @@ use crate::{
     shard_set::{Shard, ShardRef, ShardSet},
     topology::ClusterTopology,
     Cluster, ClusterNodesNotFound, ClusterNodesResp, EtcdClientFailureWithCause, InvalidArguments,
-    MetaClientFailure, OpenShard, OpenShardWithCause, Result, ShardNotFound,
+    MetaClientFailure, OpenShard, Result, ShardNotFound,
 };
 
 /// ClusterImpl is an implementation of [`Cluster`] based [`MetaClient`].
@@ -249,18 +248,33 @@ impl Inner {
             );
         }
 
-        let req = GetTablesOfShardsRequest {
-            shard_ids: vec![shard_info.id],
-        };
+        let resp = future_ext::retry_async(
+            || {
+                let req = GetTablesOfShardsRequest {
+                    shard_ids: vec![shard_info.id],
+                };
 
-        let mut resp = self
-            .meta_client
-            .get_tables_of_shards(req)
-            .await
-            .box_err()
-            .with_context(|| OpenShardWithCause {
-                msg: format!("shard_info:{shard_info:?}"),
-            })?;
+                async {
+                    self.meta_client
+                        .get_tables_of_shards(req)
+                        .await
+                        .map_err(|e| {
+                            error!("Get tables of shard failed, err:{e}");
+                            e
+                        })
+                }
+            },
+            // TODO: configure retry
+            &future_ext::RetryConfig {
+                max_retries: 10,
+                interval: Duration::from_secs(5),
+            },
+        )
+        .await;
+        let mut resp = match resp {
+            Ok(v) => v,
+            Err(_e) => panic!("Release shard lock failed and we have to panic otherwise this shard wont be assigned again."),
+        };
 
         ensure!(
             resp.tables_by_shard.len() == 1,

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,95 @@
+// Copyright 2022-2023 CeresDB Project Authors. Licensed under Apache-2.0.
+
+//! Util function to retry future.
+
+use std::time::Duration;
+
+use futures::Future;
+
+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 - 1 {
+        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!(3, 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;
@@ -362,10 +365,24 @@ async fn do_close_shard(ctx: &HandlerContext, shard_id: ShardId) -> Result<()> {
             msg: "fail to close shards in cluster",
         })?;
 
-    ctx.release_shard_lock(shard_id).await.map_err(|e| {
-        error!("Failed to release shard lock, shard_id:{shard_id}, err:{e}");
-        e
-    })?;
+    if future_ext::retry_async(
+        || async {
+            ctx.release_shard_lock(shard_id).await.map_err(|e| {
+                error!("Failed to release shard lock, shard_id:{shard_id}, err:{e}");
+                e
+            })
+        },
+        // TODO: configure retry
+        &future_ext::RetryConfig {
+            max_retries: 10,
+            interval: Duration::from_secs(5),
+        },
+    )
+    .await
+    .is_err()
+    {
+        panic!("Release shard lock failed and we have to panic otherwise this shard wont be assigned again.")
+    }
 
     info!("Shard close sequentially finish, shard_id:{shard_id}");