From 7a9608c5a8b8c8ea5ca36c5c2af235b16d6d4ce0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Artur=20Wyszy=C5=84ski?=
 <artur.wyszynski@wildpixelgames.com>
Date: Fri, 6 Dec 2024 13:01:46 +0100
Subject: [PATCH] fix: use ConnectionGuard to track TCP connections correctly

- Add ConnectionGuard in handle_client to track active connections
- Fix connection statistics not being updated
- Add debug logging for connection lifecycle
---
 src/connection/guard.rs | 7 +++++++
 src/modbus_relay.rs     | 6 ++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/connection/guard.rs b/src/connection/guard.rs
index ec1f4f4..1e4e8d0 100644
--- a/src/connection/guard.rs
+++ b/src/connection/guard.rs
@@ -1,6 +1,7 @@
 use std::{net::SocketAddr, sync::Arc};
 
 use tokio::sync::OwnedSemaphorePermit;
+use tracing::debug;
 
 use super::ConnectionManager;
 
@@ -18,10 +19,16 @@ impl Drop for ConnectionGuard {
         let manager = Arc::clone(&self.manager);
         let addr = self.addr;
 
+        debug!("Closing connection from {}", addr);
+
         tokio::spawn(async move {
             let mut stats = manager.stats.lock().await;
             if let Some(client_stats) = stats.get_mut(&addr) {
                 client_stats.active_connections -= 1;
+                debug!(
+                    "Connection from {} closed, active connections: {}",
+                    addr, client_stats.active_connections
+                );
             }
         });
     }
diff --git a/src/modbus_relay.rs b/src/modbus_relay.rs
index bb4b54e..4352690 100644
--- a/src/modbus_relay.rs
+++ b/src/modbus_relay.rs
@@ -467,12 +467,14 @@ async fn handle_client(
     manager: Arc<ConnectionManager>,
 ) -> Result<(), RelayError> {
     let start_time = Instant::now();
-    manager.record_request(peer_addr, true).await;
+
+    // Create connection guard to track this connection
+    let _guard = manager.accept_connection(peer_addr).await?;
 
     let result = handle_client_inner(stream, peer_addr, transport, manager.clone()).await;
 
     if result.is_err() {
-        manager.record_request(peer_addr, false).await;
+        manager.record_client_error(&peer_addr).await?;
     }
 
     manager.record_response_time(start_time.elapsed()).await;