g3proxy: reuse socks5 udp wrapper stats

bytedance · Aug 16, 2023 · 90d5ed3 · 90d5ed3
1 parent 1e42915
commit 90d5ed3
Show file tree

Hide file tree

Showing 10 changed files with 128 additions and 225 deletions.
diff --git a/g3proxy/src/escape/proxy_float/peer/socks5/udp_connect/mod.rs b/g3proxy/src/escape/proxy_float/peer/socks5/udp_connect/mod.rs
@@ -14,21 +14,21 @@
  * limitations under the License.
  */
 
+use std::sync::Arc;
+
 use g3_io_ext::{LimitedUdpRecv, LimitedUdpSend};
 
-use super::{NextProxyPeerInternal, ProxyFloatEscaperStats, ProxyFloatSocks5Peer};
+use super::{NextProxyPeerInternal, ProxyFloatSocks5Peer};
 use crate::escape::proxy_socks5::udp_connect::{
     ProxySocks5UdpConnectRemoteRecv, ProxySocks5UdpConnectRemoteSend,
+    ProxySocks5UdpConnectRemoteStats,
 };
 use crate::module::tcp_connect::TcpConnectTaskNotes;
 use crate::module::udp_connect::{
     ArcUdpConnectTaskRemoteStats, UdpConnectError, UdpConnectResult, UdpConnectTaskNotes,
 };
 use crate::serve::ServerTaskNotes;
 
-mod stats;
-use stats::ProxySocks5UdpConnectRemoteStats;
-
 impl ProxyFloatSocks5Peer {
     pub(super) async fn udp_connect_to<'a>(
         &'a self,
@@ -50,8 +50,10 @@ impl ProxyFloatSocks5Peer {
         udp_notes.local = Some(udp_local_addr);
         udp_notes.next = Some(udp_peer_addr);
 
-        let mut wrapper_stats =
-            ProxySocks5UdpConnectRemoteStats::new(&self.escaper_stats, task_stats);
+        let mut wrapper_stats = ProxySocks5UdpConnectRemoteStats::new(
+            Arc::clone(&self.escaper_stats) as ArcUdpConnectTaskRemoteStats,
+            task_stats,
+        );
         wrapper_stats.push_user_io_stats(self.fetch_user_upstream_io_stats(task_notes));
         let (ups_r_stats, ups_w_stats) = wrapper_stats.into_pair();
 

diff --git a/g3proxy/src/escape/proxy_float/peer/socks5/udp_connect/stats.rs b/g3proxy/src/escape/proxy_float/peer/socks5/udp_connect/stats.rs
diff --git a/g3proxy/src/escape/proxy_float/peer/socks5/udp_relay/mod.rs b/g3proxy/src/escape/proxy_float/peer/socks5/udp_relay/mod.rs
@@ -14,21 +14,20 @@
  * limitations under the License.
  */
 
+use std::sync::Arc;
+
 use g3_io_ext::{LimitedUdpRecv, LimitedUdpSend};
 
-use super::{NextProxyPeerInternal, ProxyFloatEscaperStats, ProxyFloatSocks5Peer};
+use super::{NextProxyPeerInternal, ProxyFloatSocks5Peer};
 use crate::escape::proxy_socks5::udp_relay::{
-    ProxySocks5UdpRelayRemoteRecv, ProxySocks5UdpRelayRemoteSend,
+    ProxySocks5UdpRelayRemoteRecv, ProxySocks5UdpRelayRemoteSend, ProxySocks5UdpRelayRemoteStats,
 };
 use crate::module::tcp_connect::TcpConnectTaskNotes;
 use crate::module::udp_relay::{
     ArcUdpRelayTaskRemoteStats, UdpRelaySetupError, UdpRelaySetupResult, UdpRelayTaskNotes,
 };
 use crate::serve::ServerTaskNotes;
 
-mod stats;
-use stats::ProxySocks5UdpRelayRemoteStats;
-
 impl ProxyFloatSocks5Peer {
     pub(super) async fn udp_setup_relay<'a>(
         &'a self,
@@ -42,8 +41,10 @@ impl ProxyFloatSocks5Peer {
             .await
             .map_err(UdpRelaySetupError::SetupSocketFailed)?;
 
-        let mut wrapper_stats =
-            ProxySocks5UdpRelayRemoteStats::new(&self.escaper_stats, task_stats);
+        let mut wrapper_stats = ProxySocks5UdpRelayRemoteStats::new(
+            Arc::clone(&self.escaper_stats) as ArcUdpRelayTaskRemoteStats,
+            task_stats,
+        );
         wrapper_stats.push_user_io_stats(self.fetch_user_upstream_io_stats(task_notes));
         let (ups_r_stats, ups_w_stats) = wrapper_stats.into_pair();
 

diff --git a/g3proxy/src/escape/proxy_float/peer/socks5/udp_relay/stats.rs b/g3proxy/src/escape/proxy_float/peer/socks5/udp_relay/stats.rs
diff --git a/g3proxy/src/escape/proxy_float/stats/escaper.rs b/g3proxy/src/escape/proxy_float/stats/escaper.rs
@@ -27,6 +27,8 @@ use g3_types::stats::{StatId, TcpIoSnapshot, UdpIoSnapshot};
 use crate::escape::{
     EscaperInterfaceStats, EscaperInternalStats, EscaperStats, EscaperTcpStats, EscaperUdpStats,
 };
+use crate::module::udp_connect::UdpConnectTaskRemoteStats;
+use crate::module::udp_relay::UdpRelayTaskRemoteStats;
 
 pub(crate) struct ProxyFloatEscaperStats {
     name: MetricsName,
@@ -123,3 +125,39 @@ impl LimitedWriterStats for ProxyFloatEscaperStats {
         self.tcp.io.add_out_bytes(size);
     }
 }
+
+impl UdpRelayTaskRemoteStats for ProxyFloatEscaperStats {
+    fn add_recv_bytes(&self, size: u64) {
+        self.udp.io.add_in_bytes(size);
+    }
+
+    fn add_recv_packet(&self) {
+        self.udp.io.add_in_packet();
+    }
+
+    fn add_send_bytes(&self, size: u64) {
+        self.udp.io.add_out_bytes(size);
+    }
+
+    fn add_send_packet(&self) {
+        self.udp.io.add_out_packet();
+    }
+}
+
+impl UdpConnectTaskRemoteStats for ProxyFloatEscaperStats {
+    fn add_recv_bytes(&self, size: u64) {
+        self.udp.io.add_in_bytes(size);
+    }
+
+    fn add_recv_packet(&self) {
+        self.udp.io.add_in_packet();
+    }
+
+    fn add_send_bytes(&self, size: u64) {
+        self.udp.io.add_out_bytes(size);
+    }
+
+    fn add_send_packet(&self) {
+        self.udp.io.add_out_packet();
+    }
+}
diff --git a/g3proxy/src/escape/proxy_socks5/stats.rs b/g3proxy/src/escape/proxy_socks5/stats.rs
@@ -27,6 +27,8 @@ use g3_types::stats::{StatId, TcpIoSnapshot, UdpIoSnapshot};
 use crate::escape::{
     EscaperInterfaceStats, EscaperInternalStats, EscaperStats, EscaperTcpStats, EscaperUdpStats,
 };
+use crate::module::udp_connect::UdpConnectTaskRemoteStats;
+use crate::module::udp_relay::UdpRelayTaskRemoteStats;
 
 pub(super) struct ProxySocks5EscaperStats {
     name: MetricsName,
@@ -123,3 +125,39 @@ impl LimitedWriterStats for ProxySocks5EscaperStats {
         self.tcp.io.add_out_bytes(size);
     }
 }
+
+impl UdpRelayTaskRemoteStats for ProxySocks5EscaperStats {
+    fn add_recv_bytes(&self, size: u64) {
+        self.udp.io.add_in_bytes(size);
+    }
+
+    fn add_recv_packet(&self) {
+        self.udp.io.add_in_packet();
+    }
+
+    fn add_send_bytes(&self, size: u64) {
+        self.udp.io.add_out_bytes(size);
+    }
+
+    fn add_send_packet(&self) {
+        self.udp.io.add_out_packet();
+    }
+}
+
+impl UdpConnectTaskRemoteStats for ProxySocks5EscaperStats {
+    fn add_recv_bytes(&self, size: u64) {
+        self.udp.io.add_in_bytes(size);
+    }
+
+    fn add_recv_packet(&self) {
+        self.udp.io.add_in_packet();
+    }
+
+    fn add_send_bytes(&self, size: u64) {
+        self.udp.io.add_out_bytes(size);
+    }
+
+    fn add_send_packet(&self) {
+        self.udp.io.add_out_packet();
+    }
+}
diff --git a/g3proxy/src/escape/proxy_socks5/udp_connect/mod.rs b/g3proxy/src/escape/proxy_socks5/udp_connect/mod.rs
@@ -15,16 +15,17 @@
  */
 
 use g3_io_ext::{LimitedUdpRecv, LimitedUdpSend};
+use std::sync::Arc;
 
-use super::{ProxySocks5Escaper, ProxySocks5EscaperStats};
+use super::ProxySocks5Escaper;
 use crate::module::tcp_connect::TcpConnectTaskNotes;
 use crate::module::udp_connect::{
     ArcUdpConnectTaskRemoteStats, UdpConnectError, UdpConnectResult, UdpConnectTaskNotes,
 };
 use crate::serve::ServerTaskNotes;
 
 mod stats;
-use stats::ProxySocks5UdpConnectRemoteStats;
+pub(crate) use stats::ProxySocks5UdpConnectRemoteStats;
 
 mod recv;
 mod send;
@@ -53,7 +54,10 @@ impl ProxySocks5Escaper {
         udp_notes.local = Some(udp_local_addr);
         udp_notes.next = Some(udp_peer_addr);
 
-        let mut wrapper_stats = ProxySocks5UdpConnectRemoteStats::new(&self.stats, task_stats);
+        let mut wrapper_stats = ProxySocks5UdpConnectRemoteStats::new(
+            Arc::clone(&self.stats) as ArcUdpConnectTaskRemoteStats,
+            task_stats,
+        );
         wrapper_stats.push_user_io_stats(self.fetch_user_upstream_io_stats(task_notes));
         let (ups_r_stats, ups_w_stats) = wrapper_stats.into_pair();