refactor(gateway): address Copilot review feedback on PR 134

- Switch to monotonic tokio::time::Instant for IdleTracker to prevent clock skew bugs.
- Evaluate idle deadline dynamically without polling loop.
- Decouple traffic counting logic entirely into Arc<AtomicU64> buffers, preventing any borrow-checker risks.
- Rename client bindings logically to active_clientX to prevent ambiguity.

Author: yinheli

src/gateway/relay_tcp.rs

@@ -7,7 +7,7 @@ use std::{
         atomic::{AtomicU64, Ordering},
     },
     task::{Context, Poll},
-    time::{Duration, Instant},
+    time::Duration,
 };
 
 use anyhow::{Context as _, Result};
@@ -109,19 +109,23 @@ async fn copy_with_idle_timeout(
     target_addr: &str,
 ) -> Result<()> {
     let tracker = Arc::new(SharedIdleTracker::new());
+    let client_written = Arc::new(AtomicU64::new(0));
+    let proxy_written = Arc::new(AtomicU64::new(0));
 
-    let mut active_client = ActiveStream::new(client, tracker.clone());
-    let mut active_proxy = ActiveStream::new(proxy, tracker.clone());
+    let mut active_client = ActiveStream::new(client, tracker.clone(), client_written.clone());
+    let mut active_proxy = ActiveStream::new(proxy, tracker.clone(), proxy_written.clone());
 
     let copy_task = copy_bidirectional(&mut active_client, &mut active_proxy);
 
     let monitor_task = async {
         loop {
-            let elapsed = tracker.elapsed();
-            if elapsed >= IDLE_TIMEOUT {
+            let last = tracker.last_activity_instant();
+            let deadline = last + IDLE_TIMEOUT;
+            tokio::time::sleep_until(deadline).await;
+            if tracker.last_activity_instant() <= last {
+                // No activity since we woke up (or last activity is older) -> idle timeout
                 break;
             }
-            tokio::time::sleep(IDLE_TIMEOUT - elapsed).await;
         }
     };
 
@@ -131,13 +135,13 @@ async fn copy_with_idle_timeout(
                 Ok((up, down)) => (up, down),
                 Err(e) => {
                     debug!("TCP relay error or closed: {}", e);
-                    (active_proxy.written_bytes, active_client.written_bytes)
+                    (proxy_written.load(Ordering::Relaxed), client_written.load(Ordering::Relaxed))
                 }
             }
         }
         _ = monitor_task => {
             debug!("TCP relay idle timeout for {}", target_addr);
-            (active_proxy.written_bytes, active_client.written_bytes)
+            (proxy_written.load(Ordering::Relaxed), client_written.load(Ordering::Relaxed))
         }
     };
 
@@ -161,42 +165,41 @@ async fn find_session_target(
 }
 
 struct SharedIdleTracker {
-    base_instant: Instant,
-    last_activity: Arc<AtomicU64>,
+    base_instant: tokio::time::Instant,
+    last_activity_millis: Arc<AtomicU64>,
 }
 
 impl SharedIdleTracker {
     fn new() -> Self {
         Self {
-            base_instant: Instant::now(),
-            last_activity: Arc::new(AtomicU64::new(0)),
+            base_instant: tokio::time::Instant::now(),
+            last_activity_millis: Arc::new(AtomicU64::new(0)),
         }
     }
 
     fn update_activity(&self) {
         let elapsed = self.base_instant.elapsed().as_millis() as u64;
-        self.last_activity.store(elapsed, Ordering::Relaxed);
+        self.last_activity_millis.store(elapsed, Ordering::Relaxed);
     }
 
-    fn elapsed(&self) -> Duration {
-        let last_millis = self.last_activity.load(Ordering::Relaxed);
-        let now_millis = self.base_instant.elapsed().as_millis() as u64;
-        Duration::from_millis(now_millis.saturating_sub(last_millis))
+    fn last_activity_instant(&self) -> tokio::time::Instant {
+        let millis = self.last_activity_millis.load(Ordering::Relaxed);
+        self.base_instant + Duration::from_millis(millis)
     }
 }
 
 struct ActiveStream<T> {
     inner: T,
     tracker: Arc<SharedIdleTracker>,
-    written_bytes: u64,
+    written_bytes: Arc<AtomicU64>,
 }
 
 impl<T> ActiveStream<T> {
-    fn new(inner: T, tracker: Arc<SharedIdleTracker>) -> Self {
+    fn new(inner: T, tracker: Arc<SharedIdleTracker>, written_bytes: Arc<AtomicU64>) -> Self {
         Self {
             inner,
             tracker,
-            written_bytes: 0,
+            written_bytes,
         }
     }
 
@@ -233,11 +236,11 @@ impl<T: AsyncWrite + Unpin> AsyncWrite for ActiveStream<T> {
     ) -> Poll<Result<usize, Error>> {
         let poll = Pin::new(&mut self.inner).poll_write(cx, buf);
 
-        if let Poll::Ready(Ok(n)) = &poll
-            && *n > 0
-        {
-            self.update_activity();
-            self.written_bytes += *n as u64;
+        if let Poll::Ready(Ok(n)) = &poll {
+            if *n > 0 {
+                self.update_activity();
+                self.written_bytes.fetch_add(*n as u64, Ordering::Relaxed);
+            }
         }
 
         poll