fix: meta-semaphore: re-connect when no event recevied (#18690)

Refactor permit implementation to properly handle acquirer closure and
connection errors.

The permit now uses oneshot channels to detect when the underlying
permit entry is removed from meta-service, replacing the previous
BoxFuture. Because a future that is not polled will block the channel.

If the meta-service does not send an event for `permit_ttl * 1.5`,  the
connection will be re-established.

Add `AcquirerClosed` error type for client side close error.

Author: drmingdrmer

src/meta/binaries/metabench/main.rs

@@ -412,14 +412,14 @@ async fn benchmark_semaphore(
 
     let permit_str = format!("({sem_key}, id={id})");
 
-    let mut sem = Semaphore::new(client.clone(), &sem_key, param.capacity).await;
+    let mut sem = Semaphore::new(client.clone(), &sem_key, param.capacity, param.ttl()).await;
     if param.time_based {
         sem.set_time_based_seq(None);
     } else {
         sem.set_storage_based_seq();
     }
 
-    let permit_res = sem.acquire(&id, param.ttl()).await;
+    let permit_res = sem.acquire(&id).await;
 
     print_sem_res(i, format!("sem-acquired: {permit_str}",), &permit_res);
 

src/meta/semaphore/src/acquirer/permit.rs

@@ -14,16 +14,16 @@
 
 use std::future::Future;
 
-use futures::future::BoxFuture;
+use databend_common_base::runtime::spawn_named;
 use futures::FutureExt;
 use log::debug;
 use log::info;
+use log::warn;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
 
 use crate::acquirer::Acquirer;
 use crate::acquirer::SharedAcquirerStat;
-use crate::errors::ConnectionClosed;
 use crate::queue::PermitEvent;
 use crate::storage::PermitEntry;
 use crate::storage::PermitKey;
@@ -43,7 +43,15 @@ use crate::storage::PermitKey;
 pub struct Permit {
     pub acquirer_name: String,
     pub stat: SharedAcquirerStat,
-    pub fu: BoxFuture<'static, Result<(), ConnectionClosed>>,
+
+    // Hold it so that the subscriber keep running.
+    pub(crate) _subscriber_cancel_tx: oneshot::Sender<()>,
+
+    // Hold it so that the lease extending task keep running.
+    pub(crate) _leaser_cancel_tx: oneshot::Sender<()>,
+
+    /// Gets ready if the [`PermitEntry`] is removed from meta-service.
+    pub(crate) is_removed_rx: oneshot::Receiver<()>,
 }
 
 impl std::fmt::Debug for Permit {
@@ -68,13 +76,16 @@ impl Drop for Permit {
 }
 
 impl Future for Permit {
-    type Output = Result<(), ConnectionClosed>;
+    type Output = ();
 
     fn poll(
         mut self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
     ) -> std::task::Poll<Self::Output> {
-        self.fu.poll_unpin(cx)
+        match self.is_removed_rx.poll_unpin(cx) {
+            std::task::Poll::Ready(_) => std::task::Poll::Ready(()),
+            std::task::Poll::Pending => std::task::Poll::Pending,
+        }
     }
 }
 
@@ -89,18 +100,26 @@ impl Permit {
         permit_entry: PermitEntry,
         leaser_cancel_tx: oneshot::Sender<()>,
     ) -> Self {
+        let (is_removed_tx, is_removed_rx) = oneshot::channel::<()>();
+
+        // There must be a standalone task that consumes incoming events so that it won't block the permit_event_rx sender
+        let acquirer_name = acquirer.name.clone();
         let fu = Self::watch_for_remove(
             acquirer.permit_event_rx,
+            acquirer_name.clone(),
             permit_key,
             permit_entry,
-            acquirer.subscriber_cancel_tx,
-            leaser_cancel_tx,
+            is_removed_tx,
         );
 
+        spawn_named(fu, format!("{}-WatchRemove", acquirer_name.clone()));
+
         Permit {
-            acquirer_name: acquirer.name,
+            acquirer_name,
             stat: acquirer.stat,
-            fu: Box::pin(fu),
+            _subscriber_cancel_tx: acquirer.subscriber_cancel_tx,
+            _leaser_cancel_tx: leaser_cancel_tx,
+            is_removed_rx,
         }
     }
 
@@ -120,30 +139,26 @@ impl Permit {
     /// in such case, the permit may still be valid.
     pub(crate) async fn watch_for_remove(
         mut permit_event_rx: mpsc::Receiver<PermitEvent>,
+        acquirer_name: String,
         permit_key: PermitKey,
         permit_entry: PermitEntry,
-        _subscriber_cancel_tx: oneshot::Sender<()>,
-        _leaser_cancel_tx: oneshot::Sender<()>,
-    ) -> Result<(), ConnectionClosed> {
-        let ctx = format!("Semaphore-Acquired: {}->{}", permit_key, permit_entry);
+        is_removed_tx: oneshot::Sender<()>,
+    ) {
+        let ctx = format!("{}: {}->{}", acquirer_name, permit_key, permit_entry);
 
         while let Some(sem_event) = permit_event_rx.recv().await {
-            debug!("semaphore event: {} received by: {}", sem_event, ctx);
+            debug!("{}: received semaphore event: {}", ctx, sem_event);
 
             match sem_event {
                 PermitEvent::Acquired((_seq, _entry)) => {}
                 PermitEvent::Removed((seq, _)) => {
                     if seq == permit_key.seq {
-                        debug!(
-                            "semaphore PermitEntry is removed: {}->{}",
-                            permit_key, permit_entry.id
-                        );
-                        return Ok(());
+                        warn!("{}: PermitEntry is removed", ctx);
+                        is_removed_tx.send(()).ok();
+                        return;
                     }
                 }
             }
         }
-
-        Err(ConnectionClosed::new_str("semaphore event stream closed").context(&ctx))
     }
 }

src/meta/semaphore/src/errors/acquirer_closed.rs

@@ -0,0 +1,102 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::fmt;
+
+/// Error indicating that a semaphore acquirer or permit has been closed.
+///
+/// This occurs when the acquirer instance becomes unavailable due to connection loss,
+/// explicit closure, permit expiration, or meta-service unavailability.
+#[derive(thiserror::Error, Debug)]
+pub struct AcquirerClosed {
+    reason: String,
+    when: Vec<String>,
+}
+
+impl AcquirerClosed {
+    /// Creates a new acquirer closed error with the specified reason.
+    pub fn new(reason: impl ToString) -> Self {
+        AcquirerClosed {
+            reason: reason.to_string(),
+            when: vec![],
+        }
+    }
+
+    /// Adds contextual information about when or where the error occurred.
+    ///
+    /// Context is displayed in the order it was added and enables method chaining.
+    pub fn context(mut self, context: impl ToString) -> Self {
+        self.when.push(context.to_string());
+        self
+    }
+}
+
+impl fmt::Display for AcquirerClosed {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "distributed-Semaphore Acquirer or Permit is closed: {}",
+            self.reason
+        )?;
+
+        if self.when.is_empty() {
+            return Ok(());
+        }
+
+        write!(f, "; when: (")?;
+
+        for (i, when) in self.when.iter().enumerate() {
+            if i > 0 {
+                write!(f, "; ")?;
+            }
+            write!(f, "{}", when)?;
+        }
+
+        write!(f, ")")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_basic_error_creation() {
+        let err = AcquirerClosed::new("Connection timeout");
+        assert_eq!(
+            err.to_string(),
+            "distributed-Semaphore Acquirer or Permit is closed: Connection timeout"
+        );
+    }
+
+    #[test]
+    fn test_error_with_context_chain() {
+        let err = AcquirerClosed::new("Network unreachable")
+            .context("while acquiring read permit")
+            .context("during database query");
+        assert_eq!(
+            err.to_string(),
+            "distributed-Semaphore Acquirer or Permit is closed: Network unreachable; when: (while acquiring read permit; during database query)"
+        );
+    }
+
+    #[test]
+    fn test_error_without_context() {
+        let err = AcquirerClosed::new("Simple error");
+        assert_eq!(
+            err.to_string(),
+            "distributed-Semaphore Acquirer or Permit is closed: Simple error"
+        );
+    }
+}

src/meta/semaphore/src/errors/mod.rs

@@ -13,11 +13,15 @@
 // limitations under the License.
 
 mod acquire_error;
+mod acquirer_closed;
 mod connection_closed;
 mod early_removed;
 mod either;
+mod processer_error;
 
 pub use acquire_error::AcquireError;
+pub use acquirer_closed::AcquirerClosed;
 pub use connection_closed::ConnectionClosed;
 pub use early_removed::EarlyRemoved;
 pub use either::Either;
+pub use processer_error::ProcessorError;

src/meta/semaphore/src/errors/processer_error.rs

@@ -0,0 +1,36 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use tonic::Status;
+
+use crate::errors::AcquirerClosed;
+use crate::errors::ConnectionClosed;
+
+/// Errors during semaphore permit processing.
+#[derive(thiserror::Error, Debug)]
+pub enum ProcessorError {
+    /// Connection to meta-service was lost.
+    #[error("ProcessorError: {0}")]
+    ConnectionClosed(#[from] ConnectionClosed),
+
+    /// Acquirer or Permit was dropped and there is no receiving end to receive the event.
+    #[error("ProcessorError: the semaphore Acquirer or Permit is dropped {0}")]
+    AcquirerClosed(#[from] AcquirerClosed),
+}
+
+impl From<Status> for ProcessorError {
+    fn from(status: Status) -> Self {
+        ProcessorError::ConnectionClosed(status.into())
+    }
+}

src/meta/semaphore/src/meta_event_subscriber/processor.rs

@@ -21,7 +21,9 @@ use tokio::sync::mpsc;
 
 #[cfg(doc)]
 use crate::acquirer::Acquirer;
+use crate::errors::AcquirerClosed;
 use crate::errors::ConnectionClosed;
+use crate::errors::ProcessorError;
 use crate::queue::PermitEvent;
 use crate::queue::SemaphoreQueue;
 use crate::PermitEntry;
@@ -34,18 +36,21 @@ pub(crate) struct Processor {
     pub(crate) queue: SemaphoreQueue,
 
     /// The channel to send the [`PermitEvent`] to the [`Acquirer`].
-    pub(crate) tx: mpsc::Sender<PermitEvent>,
+    ///
+    /// When the [`Acquirer`] acquired a [`Permit`], this channel is no longer used,
+    /// because the receiving end may not be polled and sending to it may block forever.
+    pub(crate) tx_to_acquirer: mpsc::Sender<PermitEvent>,
 
     /// Contains descriptive information about the context of this watcher.
     pub(crate) watcher_name: String,
 }
 
 impl Processor {
     /// Create a new [`Processor`] instance with given permits capacity and the channel to send the [`PermitEvent`] to the [`Acquirer`].
-    pub(crate) fn new(capacity: u64, tx: mpsc::Sender<PermitEvent>) -> Self {
+    pub(crate) fn new(capacity: u64, tx_to_acquirer: mpsc::Sender<PermitEvent>) -> Self {
         Self {
             queue: SemaphoreQueue::new(capacity),
-            tx,
+            tx_to_acquirer,
             watcher_name: "".to_string(),
         }
     }
@@ -60,22 +65,24 @@ impl Processor {
     pub(crate) async fn process_watch_response(
         &mut self,
         watch_response: WatchResponse,
-    ) -> Result<(), ConnectionClosed> {
+    ) -> Result<(), ProcessorError> {
         let Some((key, prev, current)) =
             Self::decode_watch_response(watch_response, &self.watcher_name)?
         else {
             return Ok(());
         };
 
-        self.process_kv_change(key, prev, current).await
+        self.process_kv_change(key, prev, current).await?;
+
+        Ok(())
     }
 
     async fn process_kv_change(
         &mut self,
         sem_key: PermitKey,
         prev: Option<SeqV<PermitEntry>>,
         current: Option<SeqV<PermitEntry>>,
-    ) -> Result<(), ConnectionClosed> {
+    ) -> Result<(), AcquirerClosed> {
         log::debug!(
             "{} processing kv change: {}: {:?} -> {:?}",
             self.watcher_name,
@@ -108,8 +115,8 @@ impl Processor {
         for event in state_changes {
             log::debug!("{} sending event: {}", self.watcher_name, event);
 
-            self.tx.send(event).await.map_err(|e| {
-                ConnectionClosed::new_str(format!("Semaphore-Watcher fail to send {}", e.0))
+            self.tx_to_acquirer.send(event).await.map_err(|e| {
+                AcquirerClosed::new(format!("Semaphore-Watcher fail to send {}", e.0))
                     .context(&self.watcher_name)
             })?;
         }