增加必要的数据记录

src/rtr/server/service.rs

@@ -11,6 +11,9 @@ use tracing::{error, warn};
 
 use crate::rtr::cache::SharedRtrCache;
 use crate::rtr::server::config::RtrServiceConfig;
+use crate::rtr::server::connection::{
+    RtrTransportConnectionCounters, RtrTransportConnectionCounts,
+};
 use crate::rtr::server::listener::RtrServer;
 use crate::rtr::server::notifier::RtrNotifier;
 use crate::rtr::server::ssh::{SshAuthMode, load_rtr_ssh_runtime_config};
@@ -22,9 +25,31 @@ pub struct RtrService {
     connection_limiter: Arc<Semaphore>,
     handshake_limiter: Arc<Semaphore>,
     active_connections: Arc<AtomicUsize>,
+    transport_connections: Arc<RtrTransportConnectionCounters>,
     config: RtrServiceConfig,
 }
 
+#[derive(Clone)]
+pub struct RtrServiceStats {
+    active_connections: Arc<AtomicUsize>,
+    transport_connections: Arc<RtrTransportConnectionCounters>,
+    max_connections: usize,
+}
+
+impl RtrServiceStats {
+    pub fn active_connections(&self) -> usize {
+        self.active_connections.load(Ordering::Relaxed)
+    }
+
+    pub fn transport_connections(&self) -> RtrTransportConnectionCounts {
+        self.transport_connections.snapshot()
+    }
+
+    pub fn max_connections(&self) -> usize {
+        self.max_connections
+    }
+}
+
 impl RtrService {
     pub fn new(cache: SharedRtrCache) -> Self {
         Self::with_config(cache, RtrServiceConfig::default())
@@ -41,6 +66,7 @@ impl RtrService {
             connection_limiter: Arc::new(Semaphore::new(config.max_connections)),
             handshake_limiter: Arc::new(Semaphore::new(config.max_concurrent_handshakes)),
             active_connections: Arc::new(AtomicUsize::new(0)),
+            transport_connections: Arc::new(RtrTransportConnectionCounters::default()),
             config,
         }
     }
@@ -61,10 +87,22 @@ impl RtrService {
         self.active_connections.load(Ordering::Relaxed)
     }
 
+    pub fn transport_connections(&self) -> RtrTransportConnectionCounts {
+        self.transport_connections.snapshot()
+    }
+
     pub fn max_connections(&self) -> usize {
         self.config.max_connections
     }
 
+    pub fn stats(&self) -> RtrServiceStats {
+        RtrServiceStats {
+            active_connections: self.active_connections.clone(),
+            transport_connections: self.transport_connections.clone(),
+            max_connections: self.config.max_connections,
+        }
+    }
+
     pub fn tcp_server(&self, bind_addr: SocketAddr) -> RtrServer {
         RtrServer::new(
             bind_addr,
@@ -74,6 +112,7 @@ impl RtrService {
             self.connection_limiter.clone(),
             self.handshake_limiter.clone(),
             self.active_connections.clone(),
+            self.transport_connections.clone(),
             self.config.clone(),
         )
     }
@@ -87,6 +126,7 @@ impl RtrService {
             self.connection_limiter.clone(),
             self.handshake_limiter.clone(),
             self.active_connections.clone(),
+            self.transport_connections.clone(),
             self.config.clone(),
         )
     }
@@ -100,6 +140,7 @@ impl RtrService {
             self.connection_limiter.clone(),
             self.handshake_limiter.clone(),
             self.active_connections.clone(),
+            self.transport_connections.clone(),
             self.config.clone(),
         )
     }