feat: add positive probe && refactor monitor metrics (#41)

Author: wayslog

Cargo.lock

@@ -42,6 +42,7 @@ crc16 = "0.4"
 futures = "0.3"
 md5 = "0.7"
 rand = "0.8"
+socket2 = "0.5"
 
 [profile.release]
 debug = true

Cargo.toml

@@ -8,6 +8,7 @@ use tokio::sync::{mpsc, oneshot};
 use tokio::task::JoinHandle;
 
 use crate::backend::client::{client_request_channel, ClientId, RequestTx};
+use crate::metrics;
 
 /// Backend node representation (host:port string).
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -51,10 +52,55 @@ where
     );
 }
 
+#[derive(Clone, Copy)]
+enum SessionKind {
+    Shared,
+    Exclusive,
+}
+
+impl SessionKind {
+    fn as_str(self) -> &'static str {
+        match self {
+            SessionKind::Shared => "shared",
+            SessionKind::Exclusive => "exclusive",
+        }
+    }
+}
+
 struct SessionHandle<T: BackendRequest> {
     tx: RequestTx<SessionCommand<T>>,
     #[allow(dead_code)]
     join: JoinHandle<()>,
+    cluster: Arc<str>,
+    backend: BackendNode,
+    kind: SessionKind,
+    active: bool,
+}
+
+impl<T: BackendRequest> SessionHandle<T> {
+    fn close(mut self) {
+        if self.active {
+            metrics::pool_session_close(
+                self.cluster.as_ref(),
+                self.backend.as_str(),
+                self.kind.as_str(),
+            );
+            self.active = false;
+        }
+    }
+}
+
+impl<T: BackendRequest> Drop for SessionHandle<T> {
+    fn drop(&mut self) {
+        if self.active {
+            metrics::pool_session_close(
+                self.cluster.as_ref(),
+                self.backend.as_str(),
+                self.kind.as_str(),
+            );
+            self.active = false;
+        }
+    }
 }
 
 struct NodeSessions<T: BackendRequest> {
@@ -117,12 +163,22 @@ impl<T: BackendRequest> ConnectionPool<T> {
             let node_sessions = guard
                 .entry(node.clone())
                 .or_insert_with(|| NodeSessions::new(self.slots_per_node));
+            metrics::pool_exclusive_idle(
+                self.cluster.as_ref(),
+                node.as_str(),
+                node_sessions.exclusive_idle.len(),
+            );
             let entry = node_sessions
                 .shared
                 .get_mut(index)
                 .expect("session index within bounds");
             let handle = entry.get_or_insert_with(|| {
-                new_session_handle(connector.clone(), node.clone(), cluster.clone())
+                new_session_handle(
+                    connector.clone(),
+                    node.clone(),
+                    cluster.clone(),
+                    SessionKind::Shared,
+                )
             });
             handle.tx.clone()
         };
@@ -140,11 +196,15 @@ impl<T: BackendRequest> ConnectionPool<T> {
         {
             let mut guard = self.sessions.write();
             if let Some(node_sessions) = guard.get_mut(&node) {
-                node_sessions.shared[index] = None;
+                if let Some(handle) = node_sessions.shared[index].take() {
+                    handle.close();
+                }
                 if node_sessions.is_empty() {
+                    metrics::pool_exclusive_idle(self.cluster.as_ref(), node.as_str(), 0);
                     guard.remove(&node);
                 }
             }
+            metrics::backend_error(&self.cluster, node.as_str(), "enqueue_failed");
             return Err(anyhow!("failed to enqueue backend request: {err}"));
         }
         Ok(response_rx)
@@ -156,11 +216,22 @@ impl<T: BackendRequest> ConnectionPool<T> {
             let node_sessions = guard
                 .entry(node.clone())
                 .or_insert_with(|| NodeSessions::new(self.slots_per_node));
-            if let Some(handle) = node_sessions.exclusive_idle.pop() {
+            let handle = if let Some(handle) = node_sessions.exclusive_idle.pop() {
                 handle
             } else {
-                new_session_handle(self.connector.clone(), node.clone(), self.cluster.clone())
-            }
+                new_session_handle(
+                    self.connector.clone(),
+                    node.clone(),
+                    self.cluster.clone(),
+                    SessionKind::Exclusive,
+                )
+            };
+            metrics::pool_exclusive_idle(
+                self.cluster.as_ref(),
+                node.as_str(),
+                node_sessions.exclusive_idle.len(),
+            );
+            handle
         };
 
         ExclusiveConnection {
@@ -186,12 +257,24 @@ fn new_session_handle<T: BackendRequest>(
     connector: Arc<dyn Connector<T>>,
     node: BackendNode,
     cluster: Arc<str>,
+    kind: SessionKind,
 ) -> SessionHandle<T> {
+    let backend = node.clone();
+    let cluster_for_metrics = cluster.clone();
     let (tx, rx) = client_request_channel();
-    let join = spawn_session(connector, node, cluster, rx);
+    let join = spawn_session(connector, backend.clone(), cluster.clone(), rx);
+    metrics::pool_session_open(
+        cluster_for_metrics.as_ref(),
+        backend.as_str(),
+        kind.as_str(),
+    );
     SessionHandle {
-        tx: tx.clone(),
+        tx,
         join,
+        cluster: cluster_for_metrics,
+        backend,
+        kind,
+        active: true,
     }
 }
 
@@ -213,21 +296,29 @@ impl<'a, T: BackendRequest> ExclusiveConnection<'a, T> {
     pub async fn send(&mut self, mut request: T) -> Result<oneshot::Receiver<Result<T::Response>>> {
         request.apply_total_tracker(&self.pool.cluster);
         request.apply_remote_tracker(&self.pool.cluster);
-        let handle = self
+        let tx = self
             .handle
             .as_ref()
-            .ok_or_else(|| anyhow!("exclusive connection has been released"))?;
+            .ok_or_else(|| anyhow!("exclusive connection has been released"))?
+            .tx
+            .clone();
 
         let (respond_to, response_rx) = oneshot::channel();
-        if let Err(err) = handle
-            .tx
+        if let Err(err) = tx
             .send(SessionCommand {
                 request,
                 respond_to,
             })
             .await
         {
-            self.handle = None;
+            if let Some(handle) = self.handle.take() {
+                handle.close();
+            }
+            metrics::backend_error(
+                &self.pool.cluster,
+                self.node.as_str(),
+                "exclusive_enqueue_failed",
+            );
             return Err(anyhow!("failed to enqueue backend request: {err}"));
         }
         Ok(response_rx)
@@ -246,6 +337,11 @@ impl<'a, T: BackendRequest> Drop for ExclusiveConnection<'a, T> {
                 .entry(self.node.clone())
                 .or_insert_with(|| NodeSessions::new(self.pool.slots_per_node));
             node_sessions.exclusive_idle.push(handle);
+            metrics::pool_exclusive_idle(
+                self.pool.cluster.as_ref(),
+                self.node.as_str(),
+                node_sessions.exclusive_idle.len(),
+            );
         }
     }
 }