Conntrack Control Method

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/proxy/shared_state.rs

@@ -1,15 +1,40 @@
 use std::collections::HashSet;
 use std::collections::hash_map::RandomState;
-use std::net::IpAddr;
-use std::sync::atomic::AtomicU64;
+use std::net::{IpAddr, SocketAddr};
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::Instant;
 
 use dashmap::DashMap;
+use tokio::sync::mpsc;
 
 use crate::proxy::handshake::{AuthProbeState, AuthProbeSaturationState};
 use crate::proxy::middle_relay::{DesyncDedupRotationState, RelayIdleCandidateRegistry};
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum ConntrackCloseReason {
+    NormalEof,
+    Timeout,
+    Pressure,
+    Reset,
+    Other,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct ConntrackCloseEvent {
+    pub(crate) src: SocketAddr,
+    pub(crate) dst: SocketAddr,
+    pub(crate) reason: ConntrackCloseReason,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum ConntrackClosePublishResult {
+    Sent,
+    Disabled,
+    QueueFull,
+    QueueClosed,
+}
+
 pub(crate) struct HandshakeSharedState {
     pub(crate) auth_probe: DashMap<IpAddr, AuthProbeState>,
     pub(crate) auth_probe_saturation: Mutex<Option<AuthProbeSaturationState>>,
@@ -31,6 +56,8 @@ pub(crate) struct MiddleRelaySharedState {
 pub(crate) struct ProxySharedState {
     pub(crate) handshake: HandshakeSharedState,
     pub(crate) middle_relay: MiddleRelaySharedState,
+    pub(crate) conntrack_pressure_active: AtomicBool,
+    pub(crate) conntrack_close_tx: Mutex<Option<mpsc::Sender<ConntrackCloseEvent>>>,
 }
 
 impl ProxySharedState {
@@ -52,6 +79,67 @@ impl ProxySharedState {
                 relay_idle_registry: Mutex::new(RelayIdleCandidateRegistry::default()),
                 relay_idle_mark_seq: AtomicU64::new(0),
             },
+            conntrack_pressure_active: AtomicBool::new(false),
+            conntrack_close_tx: Mutex::new(None),
         })
     }
+
+    pub(crate) fn set_conntrack_close_sender(&self, tx: mpsc::Sender<ConntrackCloseEvent>) {
+        match self.conntrack_close_tx.lock() {
+            Ok(mut guard) => {
+                *guard = Some(tx);
+            }
+            Err(poisoned) => {
+                let mut guard = poisoned.into_inner();
+                *guard = Some(tx);
+                self.conntrack_close_tx.clear_poison();
+            }
+        }
+    }
+
+    pub(crate) fn disable_conntrack_close_sender(&self) {
+        match self.conntrack_close_tx.lock() {
+            Ok(mut guard) => {
+                *guard = None;
+            }
+            Err(poisoned) => {
+                let mut guard = poisoned.into_inner();
+                *guard = None;
+                self.conntrack_close_tx.clear_poison();
+            }
+        }
+    }
+
+    pub(crate) fn publish_conntrack_close_event(
+        &self,
+        event: ConntrackCloseEvent,
+    ) -> ConntrackClosePublishResult {
+        let tx = match self.conntrack_close_tx.lock() {
+            Ok(guard) => guard.clone(),
+            Err(poisoned) => {
+                let guard = poisoned.into_inner();
+                let cloned = guard.clone();
+                self.conntrack_close_tx.clear_poison();
+                cloned
+            }
+        };
+
+        let Some(tx) = tx else {
+            return ConntrackClosePublishResult::Disabled;
+        };
+
+        match tx.try_send(event) {
+            Ok(()) => ConntrackClosePublishResult::Sent,
+            Err(mpsc::error::TrySendError::Full(_)) => ConntrackClosePublishResult::QueueFull,
+            Err(mpsc::error::TrySendError::Closed(_)) => ConntrackClosePublishResult::QueueClosed,
+        }
+    }
+
+    pub(crate) fn set_conntrack_pressure_active(&self, active: bool) {
+        self.conntrack_pressure_active.store(active, Ordering::Relaxed);
+    }
+
+    pub(crate) fn conntrack_pressure_active(&self) -> bool {
+        self.conntrack_pressure_active.load(Ordering::Relaxed)
+    }
 }