From ccc51875fe27610d126cf85555576f160e08999b Mon Sep 17 00:00:00 2001
From: Ximon Eighteen <3304436+ximon18@users.noreply.github.com>
Date: Thu, 4 Apr 2024 10:32:29 +0200
Subject: [PATCH] Simplify connection count checking by reporting zero
 connections for connection-less oriented server metrics.

---
 src/net/server/metrics.rs | 5 +++--
 src/net/server/stream.rs  | 5 +----
 src/net/server/tests.rs   | 2 +-
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/net/server/metrics.rs b/src/net/server/metrics.rs
index 905c068cb..47243d103 100644
--- a/src/net/server/metrics.rs
+++ b/src/net/server/metrics.rs
@@ -49,13 +49,14 @@ impl ServerMetrics {
 impl ServerMetrics {
     /// The number of current connections, if applicable.
     ///
-    /// This will be None for connection-less servers such as [`DgramServer`].
+    /// This will be zero for connection-less servers such as [`DgramServer`].
     ///
     /// [`DgramServer`]: crate::net::server::dgram::DgramServer
-    pub fn num_connections(&self) -> Option<usize> {
+    pub fn num_connections(&self) -> usize {
         self.num_connections
             .as_ref()
             .map(|atomic| atomic.load(Ordering::Relaxed))
+            .unwrap_or_default()
     }
 
     /// Set the number of current connections metric.
diff --git a/src/net/server/stream.rs b/src/net/server/stream.rs
index 64b109e3b..797f7b899 100644
--- a/src/net/server/stream.rs
+++ b/src/net/server/stream.rs
@@ -523,10 +523,7 @@ where
                 accept_res = self.accept() => {
                     match accept_res {
                         Ok((stream, addr)) => {
-                            // SAFETY: This is a connection-oriented server so there
-                            // must always be a connection count metric avasilable to
-                            // unwrap.
-                            let num_conn = self.metrics.num_connections().unwrap();
+                            let num_conn = self.metrics.num_connections();
                             if num_conn < self.config.load().max_concurrent_connections {
                                 self.spawn_connection_handler(stream, addr);
                             }
diff --git a/src/net/server/tests.rs b/src/net/server/tests.rs
index ae42529e9..65586d41a 100644
--- a/src/net/server/tests.rs
+++ b/src/net/server/tests.rs
@@ -414,7 +414,7 @@ async fn service_test() {
 
         // Verify that no requests or responses are in progress still in
         // the server.
-        assert_eq!(srv.metrics().num_connections(), Some(0));
+        assert_eq!(srv.metrics().num_connections(), 0);
         assert_eq!(srv.metrics().num_inflight_requests(), 0);
         assert_eq!(srv.metrics().num_pending_writes(), 0);
         assert_eq!(srv.metrics().num_received_requests(), num_messages);