smartproxy/rust/crates/rustproxy-passthrough/src/connection_record.rs

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::{Duration, Instant};

/// Per-connection tracking record with atomics for lock-free updates.
///
/// Each field uses atomics so that the forwarding tasks can update
/// bytes_received / bytes_sent / last_activity without holding any lock,
/// while the zombie scanner reads them concurrently.
pub struct ConnectionRecord {
    /// Unique connection ID assigned by the ConnectionTracker.
    pub id: u64,
    /// Wall-clock instant when this connection was created.
    pub created_at: Instant,
    /// Milliseconds since `created_at` when the last activity occurred.
    /// Updated atomically by the forwarding loops.
    pub last_activity: AtomicU64,
    /// Total bytes received from the client (inbound).
    pub bytes_received: AtomicU64,
    /// Total bytes sent to the client (outbound / from backend).
    pub bytes_sent: AtomicU64,
    /// True once the client side of the connection has closed.
    pub client_closed: AtomicBool,
    /// True once the backend side of the connection has closed.
    pub backend_closed: AtomicBool,
    /// Whether this connection uses TLS (affects zombie thresholds).
    pub is_tls: AtomicBool,
    /// Whether this connection has keep-alive semantics.
    pub has_keep_alive: AtomicBool,
}

impl ConnectionRecord {
    /// Create a new connection record with the given ID.
    /// All counters start at zero, all flags start as false.
    pub fn new(id: u64) -> Self {
        Self {
            id,
            created_at: Instant::now(),
            last_activity: AtomicU64::new(0),
            bytes_received: AtomicU64::new(0),
            bytes_sent: AtomicU64::new(0),
            client_closed: AtomicBool::new(false),
            backend_closed: AtomicBool::new(false),
            is_tls: AtomicBool::new(false),
            has_keep_alive: AtomicBool::new(false),
        }
    }

    /// Update `last_activity` to reflect the current elapsed time.
    pub fn touch(&self) {
        let elapsed_ms = self.created_at.elapsed().as_millis() as u64;
        self.last_activity.store(elapsed_ms, Ordering::Relaxed);
    }

    /// Record `n` bytes received from the client (inbound).
    pub fn record_bytes_in(&self, n: u64) {
        self.bytes_received.fetch_add(n, Ordering::Relaxed);
        self.touch();
    }

    /// Record `n` bytes sent to the client (outbound / from backend).
    pub fn record_bytes_out(&self, n: u64) {
        self.bytes_sent.fetch_add(n, Ordering::Relaxed);
        self.touch();
    }

    /// How long since the last activity on this connection.
    pub fn idle_duration(&self) -> Duration {
        let last_ms = self.last_activity.load(Ordering::Relaxed);
        let age_ms = self.created_at.elapsed().as_millis() as u64;
        Duration::from_millis(age_ms.saturating_sub(last_ms))
    }

    /// Total age of this connection (time since creation).
    pub fn age(&self) -> Duration {
        self.created_at.elapsed()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;

    #[test]
    fn test_new_record() {
        let record = ConnectionRecord::new(42);
        assert_eq!(record.id, 42);
        assert_eq!(record.bytes_received.load(Ordering::Relaxed), 0);
        assert_eq!(record.bytes_sent.load(Ordering::Relaxed), 0);
        assert!(!record.client_closed.load(Ordering::Relaxed));
        assert!(!record.backend_closed.load(Ordering::Relaxed));
        assert!(!record.is_tls.load(Ordering::Relaxed));
        assert!(!record.has_keep_alive.load(Ordering::Relaxed));
    }

    #[test]
    fn test_record_bytes() {
        let record = ConnectionRecord::new(1);
        record.record_bytes_in(100);
        record.record_bytes_in(200);
        assert_eq!(record.bytes_received.load(Ordering::Relaxed), 300);

        record.record_bytes_out(50);
        record.record_bytes_out(75);
        assert_eq!(record.bytes_sent.load(Ordering::Relaxed), 125);
    }

    #[test]
    fn test_touch_updates_activity() {
        let record = ConnectionRecord::new(1);
        assert_eq!(record.last_activity.load(Ordering::Relaxed), 0);

        // Sleep briefly so elapsed time is nonzero
        thread::sleep(Duration::from_millis(10));
        record.touch();

        let activity = record.last_activity.load(Ordering::Relaxed);
        assert!(activity >= 10, "last_activity should be at least 10ms, got {}", activity);
    }

    #[test]
    fn test_idle_duration() {
        let record = ConnectionRecord::new(1);
        // Initially idle_duration ~ age since last_activity is 0
        thread::sleep(Duration::from_millis(20));
        let idle = record.idle_duration();
        assert!(idle >= Duration::from_millis(20));

        // After touch, idle should be near zero
        record.touch();
        let idle = record.idle_duration();
        assert!(idle < Duration::from_millis(10));
    }

    #[test]
    fn test_age() {
        let record = ConnectionRecord::new(1);
        thread::sleep(Duration::from_millis(20));
        let age = record.age();
        assert!(age >= Duration::from_millis(20));
    }

    #[test]
    fn test_flags() {
        let record = ConnectionRecord::new(1);
        record.client_closed.store(true, Ordering::Relaxed);
        record.is_tls.store(true, Ordering::Relaxed);
        record.has_keep_alive.store(true, Ordering::Relaxed);

        assert!(record.client_closed.load(Ordering::Relaxed));
        assert!(!record.backend_closed.load(Ordering::Relaxed));
        assert!(record.is_tls.load(Ordering::Relaxed));
        assert!(record.has_keep_alive.load(Ordering::Relaxed));
    }
}