use Arc string value in EgressInfo

bytedance · Jun 14, 2024 · f156b6c · f156b6c
1 parent 4e47013
commit f156b6c
Show file tree

Hide file tree

Showing 8 changed files with 58 additions and 56 deletions.
diff --git a/g3proxy/src/config/escaper/direct_float/bind/json.rs b/g3proxy/src/config/escaper/direct_float/bind/json.rs
@@ -64,19 +64,19 @@ impl DirectFloatBindIp {
                         }
                         CONFIG_KEY_ISP => {
                             if let Ok(isp) = g3_json::value::as_string(v) {
-                                bind.egress_info.isp = Some(isp);
+                                bind.egress_info.set_isp(isp);
                             }
                             // not a required field, skip if value format is invalid
                         }
                         CONFIG_KEY_EIP => {
                             if let Ok(ip) = g3_json::value::as_ipaddr(v) {
-                                bind.egress_info.ip = Some(ip);
+                                bind.egress_info.set_ip(ip);
                             }
                             // not a required field, skip if value format is invalid
                         }
                         CONFIG_KEY_AREA => {
                             if let Ok(area) = g3_json::value::as_egress_area(v) {
-                                bind.egress_info.area = Some(area);
+                                bind.egress_info.set_area(area);
                             }
                             // not a required field, skip if value format is invalid
                         }

diff --git a/g3proxy/src/escape/proxy_float/peer/http/mod.rs b/g3proxy/src/escape/proxy_float/peer/http/mod.rs
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-use std::net::{IpAddr, SocketAddr};
+use std::net::SocketAddr;
 use std::sync::Arc;
 
 use anyhow::{anyhow, Context};
@@ -25,7 +25,7 @@ use tokio::time::Instant;
 
 use g3_daemon::stat::remote::ArcTcpConnectionTaskRemoteStats;
 use g3_types::auth::{Password, Username};
-use g3_types::net::{EgressArea, EgressInfo, Host, OpensslClientConfig, TcpSockSpeedLimitConfig};
+use g3_types::net::{EgressInfo, Host, OpensslClientConfig, TcpSockSpeedLimitConfig};
 
 use super::{
     ArcNextProxyPeer, NextProxyPeer, NextProxyPeerInternal, ProxyFloatEscaper,
@@ -89,16 +89,8 @@ impl ProxyFloatHttpPeer {
 }
 
 impl NextProxyPeerInternal for ProxyFloatHttpPeer {
-    fn set_isp(&mut self, isp: String) {
-        self.egress_info.isp = Some(isp);
-    }
-
-    fn set_eip(&mut self, eip: IpAddr) {
-        self.egress_info.ip = Some(eip);
-    }
-
-    fn set_area(&mut self, area: EgressArea) {
-        self.egress_info.area = Some(area);
+    fn egress_info_mut(&mut self) -> &mut EgressInfo {
+        &mut self.egress_info
     }
 
     fn set_expire(&mut self, expire_datetime: DateTime<Utc>, expire_instant: Instant) {

diff --git a/g3proxy/src/escape/proxy_float/peer/https/mod.rs b/g3proxy/src/escape/proxy_float/peer/https/mod.rs
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-use std::net::{IpAddr, SocketAddr};
+use std::net::SocketAddr;
 use std::sync::Arc;
 
 use anyhow::{anyhow, Context};
@@ -25,7 +25,7 @@ use tokio::time::Instant;
 
 use g3_daemon::stat::remote::ArcTcpConnectionTaskRemoteStats;
 use g3_types::auth::{Password, Username};
-use g3_types::net::{EgressArea, EgressInfo, Host, OpensslClientConfig, TcpSockSpeedLimitConfig};
+use g3_types::net::{EgressInfo, Host, OpensslClientConfig, TcpSockSpeedLimitConfig};
 
 use super::{ArcNextProxyPeer, NextProxyPeer, NextProxyPeerInternal, ProxyFloatEscaper};
 use crate::module::http_forward::{ArcHttpForwardTaskRemoteStats, BoxHttpForwardConnection};
@@ -89,16 +89,8 @@ impl ProxyFloatHttpsPeer {
 }
 
 impl NextProxyPeerInternal for ProxyFloatHttpsPeer {
-    fn set_isp(&mut self, isp: String) {
-        self.egress_info.isp = Some(isp);
-    }
-
-    fn set_eip(&mut self, eip: IpAddr) {
-        self.egress_info.ip = Some(eip);
-    }
-
-    fn set_area(&mut self, area: EgressArea) {
-        self.egress_info.area = Some(area);
+    fn egress_info_mut(&mut self) -> &mut EgressInfo {
+        &mut self.egress_info
     }
 
     fn set_expire(&mut self, expire_datetime: DateTime<Utc>, expire_instant: Instant) {

diff --git a/g3proxy/src/escape/proxy_float/peer/json.rs b/g3proxy/src/escape/proxy_float/peer/json.rs
@@ -57,19 +57,19 @@ pub(super) fn do_parse_peer(
                 }
                 CONFIG_KEY_PEER_ISP => {
                     if let Ok(isp) = g3_json::value::as_string(v) {
-                        peer_mut.set_isp(isp);
+                        peer_mut.egress_info_mut().set_isp(isp);
                     }
                     // not a required field, skip if value format is invalid
                 }
                 CONFIG_KEY_PEER_EIP => {
                     if let Ok(ip) = g3_json::value::as_ipaddr(v) {
-                        peer_mut.set_eip(ip);
+                        peer_mut.egress_info_mut().set_ip(ip);
                     }
                     // not a required field, skip if value format is invalid
                 }
                 CONFIG_KEY_PEER_AREA => {
                     if let Ok(area) = g3_json::value::as_egress_area(v) {
-                        peer_mut.set_area(area);
+                        peer_mut.egress_info_mut().set_area(area);
                     }
                     // not a required field, skip if value format is invalid
                 }

diff --git a/g3proxy/src/escape/proxy_float/peer/mod.rs b/g3proxy/src/escape/proxy_float/peer/mod.rs
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-use std::net::{IpAddr, SocketAddr};
+use std::net::SocketAddr;
 use std::sync::Arc;
 
 use ahash::AHashMap;
@@ -26,7 +26,7 @@ use serde_json::Value;
 use tokio::time::Instant;
 
 use g3_daemon::stat::remote::ArcTcpConnectionTaskRemoteStats;
-use g3_types::net::{EgressArea, EgressInfo, Host, OpensslClientConfig, TcpSockSpeedLimitConfig};
+use g3_types::net::{EgressInfo, Host, OpensslClientConfig, TcpSockSpeedLimitConfig};
 
 use super::{ProxyFloatEscaper, ProxyFloatEscaperConfig, ProxyFloatEscaperStats};
 use crate::module::http_forward::{ArcHttpForwardTaskRemoteStats, BoxHttpForwardConnection};
@@ -55,9 +55,7 @@ const CONFIG_KEY_PEER_AREA: &str = "area";
 const CONFIG_KEY_PEER_TCP_SOCK_SPEED_LIMIT: &str = "tcp_sock_speed_limit";
 
 pub(super) trait NextProxyPeerInternal {
-    fn set_isp(&mut self, isp: String);
-    fn set_eip(&mut self, eip: IpAddr);
-    fn set_area(&mut self, area: EgressArea);
+    fn egress_info_mut(&mut self) -> &mut EgressInfo;
     fn set_expire(&mut self, expire_datetime: DateTime<Utc>, expire_instant: Instant);
     fn set_tcp_sock_speed_limit(&mut self, speed_limit: TcpSockSpeedLimitConfig);
     fn set_kv(&mut self, k: &str, v: &Value) -> anyhow::Result<()>;

diff --git a/g3proxy/src/escape/proxy_float/peer/socks5/mod.rs b/g3proxy/src/escape/proxy_float/peer/socks5/mod.rs
@@ -28,7 +28,7 @@ use tokio::time::Instant;
 use g3_daemon::stat::remote::ArcTcpConnectionTaskRemoteStats;
 use g3_types::auth::{Password, Username};
 use g3_types::net::{
-    EgressArea, EgressInfo, Host, OpensslClientConfig, SocksAuth, TcpSockSpeedLimitConfig,
+    EgressInfo, Host, OpensslClientConfig, SocksAuth, TcpSockSpeedLimitConfig,
     UdpSockSpeedLimitConfig,
 };
 
@@ -116,16 +116,8 @@ impl ProxyFloatSocks5Peer {
 }
 
 impl NextProxyPeerInternal for ProxyFloatSocks5Peer {
-    fn set_isp(&mut self, isp: String) {
-        self.egress_info.isp = Some(isp);
-    }
-
-    fn set_eip(&mut self, eip: IpAddr) {
-        self.egress_info.ip = Some(eip);
-    }
-
-    fn set_area(&mut self, area: EgressArea) {
-        self.egress_info.area = Some(area);
+    fn egress_info_mut(&mut self) -> &mut EgressInfo {
+        &mut self.egress_info
     }
 
     fn set_expire(&mut self, expire_datetime: DateTime<Utc>, expire_instant: Instant) {

diff --git a/g3proxy/src/module/http_header/custom.rs b/g3proxy/src/module/http_header/custom.rs
@@ -102,13 +102,13 @@ fn set_value_for_dynamic_egress_info(
     egress: &EgressInfo,
 ) {
     v.extend_from_slice(server_id.as_bytes());
-    if let Some(isp) = &egress.isp {
+    if let Some(isp) = &egress.isp() {
         let _ = write!(v, "; isp={}", BASE64_STANDARD.encode(isp));
     }
-    if let Some(ip) = &egress.ip {
+    if let Some(ip) = &egress.ip() {
         let _ = write!(v, "; ip={ip}");
     }
-    if let Some(area) = &egress.area {
+    if let Some(area) = &egress.area() {
         let _ = write!(v, "; area={}", BASE64_STANDARD.encode(area.to_string()));
     }
 }

diff --git a/lib/g3-types/src/net/egress.rs b/lib/g3-types/src/net/egress.rs
@@ -17,10 +17,11 @@
 use std::fmt;
 use std::net::IpAddr;
 use std::str::FromStr;
+use std::sync::Arc;
 
 #[derive(Clone, Debug)]
 pub struct EgressArea {
-    inner: Vec<String>,
+    inner: Vec<Arc<str>>,
 }
 
 impl fmt::Display for EgressArea {
@@ -34,12 +35,12 @@ impl FromStr for EgressArea {
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         let raw: Vec<_> = s.trim().split('/').collect();
-        let mut inner = Vec::<String>::with_capacity(raw.len());
-        for s in raw.iter() {
+        let mut inner = Vec::with_capacity(raw.len());
+        for s in raw {
             if s.is_empty() {
                 continue;
             }
-            inner.push(s.to_string());
+            inner.push(Arc::from(s));
         }
         if inner.is_empty() {
             return Err(());
@@ -50,9 +51,9 @@ impl FromStr for EgressArea {
 
 #[derive(Clone, Debug, Default)]
 pub struct EgressInfo {
-    pub ip: Option<IpAddr>,
-    pub isp: Option<String>,
-    pub area: Option<EgressArea>,
+    ip: Option<IpAddr>,
+    isp: Option<Arc<str>>,
+    area: Option<EgressArea>,
 }
 
 impl EgressInfo {
@@ -61,4 +62,31 @@ impl EgressInfo {
         self.isp = None;
         self.area = None;
     }
+
+    #[inline]
+    pub fn ip(&self) -> Option<IpAddr> {
+        self.ip
+    }
+
+    pub fn set_ip(&mut self, ip: IpAddr) {
+        self.ip = Some(ip);
+    }
+
+    #[inline]
+    pub fn isp(&self) -> Option<&str> {
+        self.isp.as_deref()
+    }
+
+    pub fn set_isp(&mut self, isp: String) {
+        self.isp = Some(Arc::from(isp));
+    }
+
+    #[inline]
+    pub fn area(&self) -> Option<&EgressArea> {
+        self.area.as_ref()
+    }
+
+    pub fn set_area(&mut self, area: EgressArea) {
+        self.area = Some(area);
+    }
 }