Clever channel timeout (#851)

* active connection liveness check for channels

* clippy

lib/api/src/grpc/transport_channel_pool.rs

@@ -5,15 +5,24 @@ use std::time::{Duration, Instant};
 use std::vec::Vec;
 
 use rand::seq::SliceRandom;
+use tokio::select;
 use tonic::transport::{Channel, Error as TonicError, Uri};
 use tonic::{Code, Status};
 
+use crate::grpc::qdrant::qdrant_client::QdrantClient;
+use crate::grpc::qdrant::HealthCheckRequest;
+
 const DEFAULT_GRPC_TIMEOUT: Duration = Duration::from_secs(10);
+const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(1);
+
+const SMART_CONNECT_TIMEOUT: Duration = Duration::from_secs(1);
 const DEFAULT_POOL_SIZE: usize = 1;
 const CHANNEL_TTL: Duration = Duration::from_secs(5);
 
 struct ChannelPool {
     channels: Vec<Channel>,
+    /// Channel for fast connectivity test
+    fast_channel: Channel,
     init_at: Instant,
 }
 
@@ -22,14 +31,22 @@ impl ChannelPool {
         uri: Uri,
         pool_size: NonZeroUsize,
         grpc_timeout: Duration,
+        connection_timeout: Duration,
     ) -> Result<Self, TonicError> {
         let mut channels = Vec::with_capacity(pool_size.into());
         for _ in 0..pool_size.into() {
-            let channel = TransportChannelPool::make_channel(grpc_timeout, uri.clone()).await?;
+            let channel =
+                TransportChannelPool::make_channel(grpc_timeout, connection_timeout, uri.clone())
+                    .await?;
             channels.push(channel);
         }
+        let fast_channel =
+            TransportChannelPool::make_channel(SMART_CONNECT_TIMEOUT, connection_timeout, uri)
+                .await?;
+
         Ok(Self {
             channels,
+            fast_channel,
             init_at: Instant::now(),
         })
     }
@@ -57,6 +74,7 @@ pub struct TransportChannelPool {
     uri_to_pool: tokio::sync::RwLock<HashMap<Uri, ChannelPool>>,
     pool_size: NonZeroUsize,
     grpc_timeout: Duration,
+    connection_timeout: Duration,
 }
 
 impl Default for TransportChannelPool {
@@ -65,23 +83,29 @@ impl Default for TransportChannelPool {
             uri_to_pool: tokio::sync::RwLock::new(HashMap::new()),
             pool_size: NonZeroUsize::new(DEFAULT_POOL_SIZE).unwrap(),
             grpc_timeout: DEFAULT_GRPC_TIMEOUT,
+            connection_timeout: DEFAULT_CONNECT_TIMEOUT,
         }
     }
 }
 
 impl TransportChannelPool {
-    pub fn new(p2p_grpc_timeout: Duration, pool_size: usize) -> Self {
+    pub fn new(p2p_grpc_timeout: Duration, connection_timeout: Duration, pool_size: usize) -> Self {
         Self {
             uri_to_pool: Default::default(),
             grpc_timeout: p2p_grpc_timeout,
+            connection_timeout,
             pool_size: NonZeroUsize::new(pool_size).unwrap(),
         }
     }
 
-    pub async fn make_channel(grpc_timeout: Duration, uri: Uri) -> Result<Channel, TonicError> {
+    pub async fn make_channel(
+        grpc_timeout: Duration,
+        connection_timeout: Duration,
+        uri: Uri,
+    ) -> Result<Channel, TonicError> {
         let endpoint = Channel::builder(uri)
             .timeout(grpc_timeout)
-            .connect_timeout(grpc_timeout)
+            .connect_timeout(connection_timeout)
             .keep_alive_while_idle(true);
         // `connect` is using the `Reconnect` network service internally to handle dropped connections
         endpoint.connect().await
@@ -93,8 +117,13 @@ impl TransportChannelPool {
         let mut guard = self.uri_to_pool.write().await;
         match guard.get(&uri) {
             None => {
-                let channels =
-                    ChannelPool::init(uri.clone(), self.pool_size, self.grpc_timeout).await?;
+                let channels = ChannelPool::init(
+                    uri.clone(),
+                    self.pool_size,
+                    self.grpc_timeout,
+                    self.connection_timeout,
+                )
+                .await?;
                 let channel = channels.choose();
                 guard.insert(uri, channels);
                 Ok(channel)
@@ -113,6 +142,11 @@ impl TransportChannelPool {
         guard.get(uri).map(|channels| channels.choose())
     }
 
+    async fn get_fast_pooled_channel(&self, uri: &Uri) -> Option<Channel> {
+        let guard = self.uri_to_pool.read().await;
+        guard.get(uri).map(|channels| channels.fast_channel.clone())
+    }
+
     async fn get_or_create_pooled_channel(&self, uri: &Uri) -> Result<Channel, TonicError> {
         match self.get_pooled_channel(uri).await {
             None => self.init_pool_for_uri(uri.clone()).await,
@@ -125,6 +159,32 @@ impl TransportChannelPool {
         guard.get(uri).map(|channels| channels.init_at)
     }
 
+    /// Checks if the channel is still alive.
+    ///
+    /// It uses duplicate "fast" chanel, equivalent ot the original, but with smaller timeout.
+    /// If it can't get healthcheck response in the timeout, it assumes the channel is dead.
+    /// And we need to drop the pool for the uri and try again.
+    /// For performance reasons, we start the check only after `SMART_CONNECT_TIMEOUT`.
+    async fn check_connectability(&self, uri: &Uri) -> Status {
+        loop {
+            tokio::time::sleep(SMART_CONNECT_TIMEOUT).await;
+            let channel = self.get_fast_pooled_channel(uri).await;
+            match channel {
+                None => return Status::unavailable("Channel dropped"),
+                Some(channel) => {
+                    let mut client = QdrantClient::new(channel);
+                    let resp = client.health_check(HealthCheckRequest {}).await;
+                    match resp {
+                        Ok(_) => {
+                            // continue watching
+                        }
+                        Err(status) => return status,
+                    }
+                }
+            }
+        }
+    }
+
     // Allows to use channel to `uri`. If there is no channels to specified uri - they will be created.
     pub async fn with_channel<T, O: Future<Output = Result<T, Status>>>(
         &self,
@@ -133,13 +193,20 @@ impl TransportChannelPool {
     ) -> Result<T, RequestError<Status>> {
         let channel = self.get_or_create_pooled_channel(uri).await?;
 
-        let result = f(channel).await;
+        let result: Result<T, Status> = select! {
+            res = f(channel) => {
+                res
+            }
+            res = self.check_connectability(uri) => {
+               Err(res)
+            }
+        };
 
         // Reconnect on failure to handle the case with domain name change.
         match result {
             Ok(res) => Ok(res),
             Err(err) => match err.code() {
-                Code::Internal | Code::Unavailable => {
+                Code::Internal | Code::Unavailable | Code::Cancelled => {
                     let channel_uptime = Instant::now().duration_since(
                         self.get_created_at(uri).await.unwrap_or_else(Instant::now),
                     );

src/consensus.rs

@@ -177,6 +177,7 @@ impl Consensus {
     ) -> anyhow::Result<()> {
         // Use dedicated transport channel for bootstrapping because of specific timeout
         let channel = TransportChannelPool::make_channel(
+            Duration::from_secs(config.bootstrap_timeout_sec),
             Duration::from_secs(config.bootstrap_timeout_sec),
             bootstrap_peer,
         )
@@ -489,9 +490,10 @@ async fn who_is(
         bootstrap_uri.ok_or_else(|| anyhow::anyhow!("No bootstrap uri supplied"))?;
     let bootstrap_timeout = Duration::from_secs(config.bootstrap_timeout_sec);
     // Use dedicated transport channel for who_is because of specific timeout
-    let channel = TransportChannelPool::make_channel(bootstrap_timeout, bootstrap_uri)
-        .await
-        .context("Failed to create timeout channel")?;
+    let channel =
+        TransportChannelPool::make_channel(bootstrap_timeout, bootstrap_timeout, bootstrap_uri)
+            .await
+            .context("Failed to create timeout channel")?;
     let mut client = RaftClient::new(channel);
     Ok(client
         .who_is(tonic::Request::new(PeerId { id: peer_id }))

src/main.rs

@@ -125,8 +125,10 @@ fn main() -> anyhow::Result<()> {
     let mut channel_service = ChannelService::default();
     if settings.cluster.enabled {
         let p2p_grpc_timeout = Duration::from_millis(settings.cluster.grpc_timeout_ms);
+        let connection_timeout = Duration::from_millis(settings.cluster.connection_timeout_ms);
         channel_service.channel_pool = Arc::new(TransportChannelPool::new(
             p2p_grpc_timeout,
+            connection_timeout,
             settings.cluster.p2p.connection_pool_size,
         ));
         channel_service.id_to_address = persistent_consensus_state.peer_address_by_id.clone();

src/settings.rs

@@ -20,6 +20,8 @@ pub struct ClusterConfig {
     pub enabled: bool, // disabled by default
     #[serde(default = "default_timeout_ms")]
     pub grpc_timeout_ms: u64,
+    #[serde(default = "default_connection_timeout_ms")]
+    pub connection_timeout_ms: u64,
     #[serde(default)]
     pub p2p: P2pConfig,
     #[serde(default)]
@@ -91,6 +93,10 @@ fn default_timeout_ms() -> u64 {
     1000 * 60
 }
 
+fn default_connection_timeout_ms() -> u64 {
+    2000
+}
+
 fn default_tick_period_ms() -> u64 {
     100
 }