去除不安全的计数器

vnt/src/core/conn.rs

@@ -27,9 +27,7 @@ use crate::nat::NatTest;
 #[cfg(feature = "integrated_tun")]
 use crate::tun_tap_device::tun_create_helper::{DeviceAdapter, TunDeviceHelper};
 use crate::tun_tap_device::vnt_device::DeviceWrite;
-use crate::util::{
+use crate::util::{Scheduler, StopManager, U64Adder, WatchU64Adder};
-    Scheduler, SingleU64Adder, StopManager, U64Adder, WatchSingleU64Adder, WatchU64Adder,
-};
 use crate::{nat, VntCallback};
 
 #[derive(Clone)]
@@ -68,7 +66,7 @@ pub struct VntInner {
     context: Arc<Mutex<Option<ChannelContext>>>,
     peer_nat_info_map: Arc<RwLock<HashMap<Ipv4Addr, NatInfo>>>,
     down_count_watcher: WatchU64Adder,
-    up_count_watcher: WatchSingleU64Adder,
+    up_count_watcher: WatchU64Adder,
     client_secret_hash: Option<[u8; 16]>,
     compressor: Compressor,
     client_cipher: Cipher,
@@ -201,14 +199,13 @@ impl VntInner {
         let (punch_sender, punch_receiver) = maintain::punch_channel();
         let peer_nat_info_map: Arc<RwLock<HashMap<Ipv4Addr, NatInfo>>> =
             Arc::new(RwLock::new(HashMap::with_capacity(16)));
-        let down_counter =
+        let down_counter = U64Adder::default();
-            U64Adder::with_capacity(config.ports.as_ref().map(|v| v.len()).unwrap_or_default() + 8);
         let down_count_watcher = down_counter.watch();
         let handshake = Handshake::new(
             #[cfg(feature = "server_encrypt")]
             rsa_cipher.clone(),
         );
-        let up_counter = SingleU64Adder::new();
+        let up_counter = U64Adder::default();
         let up_count_watcher = up_counter.watch();
         #[cfg(feature = "integrated_tun")]
         let tun_device_helper = {
@@ -348,7 +345,7 @@ pub fn start<Call: VntCallback>(
     punch: Punch,
     callback: Call,
     down_count_watcher: WatchU64Adder,
-    up_count_watcher: WatchSingleU64Adder,
+    up_count_watcher: WatchU64Adder,
 ) {
     // 定时心跳
     maintain::heartbeat(

vnt/src/handle/maintain/up_status.rs

@@ -3,7 +3,7 @@ use crate::handle::CurrentDeviceInfo;
 use crate::proto::message::{ClientStatusInfo, PunchNatType, RouteItem};
 use crate::protocol::body::ENCRYPTION_RESERVED;
 use crate::protocol::{service_packet, NetPacket, Protocol, HEAD_LEN, MAX_TTL};
-use crate::util::{Scheduler, WatchSingleU64Adder, WatchU64Adder};
+use crate::util::{Scheduler, WatchU64Adder};
 use crossbeam_utils::atomic::AtomicCell;
 use protobuf::Message;
 use std::io;
@@ -16,7 +16,7 @@ pub fn up_status(
     context: ChannelContext,
     current_device_info: Arc<AtomicCell<CurrentDeviceInfo>>,
     down_count_watcher: WatchU64Adder,
-    up_count_watcher: WatchSingleU64Adder,
+    up_count_watcher: WatchU64Adder,
 ) {
     let _ = scheduler.timeout(Duration::from_secs(60), move |x| {
         up_status0(
@@ -34,7 +34,7 @@ fn up_status0(
     context: ChannelContext,
     current_device_info: Arc<AtomicCell<CurrentDeviceInfo>>,
     down_count_watcher: WatchU64Adder,
-    up_count_watcher: WatchSingleU64Adder,
+    up_count_watcher: WatchU64Adder,
 ) {
     if let Err(e) = send_up_status_packet(
         &context,
@@ -62,7 +62,7 @@ fn send_up_status_packet(
     context: &ChannelContext,
     current_device_info: &AtomicCell<CurrentDeviceInfo>,
     down_count_watcher: &WatchU64Adder,
-    up_count_watcher: &WatchSingleU64Adder,
+    up_count_watcher: &WatchU64Adder,
 ) -> io::Result<()> {
     let device_info = current_device_info.load();
     if device_info.status.offline() {

vnt/src/handle/tun_tap/tun_handler.rs

@@ -26,7 +26,7 @@ use crate::protocol;
 use crate::protocol::body::ENCRYPTION_RESERVED;
 use crate::protocol::ip_turn_packet::BroadcastPacket;
 use crate::protocol::{ip_turn_packet, NetPacket, MAX_TTL};
-use crate::util::{SingleU64Adder, StopManager};
+use crate::util::{StopManager, U64Adder};
 
 fn icmp(device_writer: &Device, mut ipv4_packet: IpV4Packet<&mut [u8]>) -> anyhow::Result<()> {
     if ipv4_packet.protocol() == Protocol::Icmp {
@@ -53,7 +53,7 @@ pub fn start(
     #[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
     client_cipher: Cipher,
     server_cipher: Cipher,
-    mut up_counter: SingleU64Adder,
+    up_counter: U64Adder,
     device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
     compressor: Compressor,
     device_stop: DeviceStop,
@@ -71,7 +71,7 @@ pub fn start(
                 ip_proxy_map,
                 client_cipher,
                 server_cipher,
-                &mut up_counter,
+                &up_counter,
                 device_list,
                 compressor,
                 device_stop,

vnt/src/handle/tun_tap/unix.rs

@@ -7,7 +7,7 @@ use crate::handle::tun_tap::DeviceStop;
 use crate::handle::{CurrentDeviceInfo, PeerDeviceInfo};
 #[cfg(feature = "ip_proxy")]
 use crate::ip_proxy::IpProxyMap;
-use crate::util::{SingleU64Adder, StopManager};
+use crate::util::{StopManager, U64Adder};
 use crossbeam_utils::atomic::AtomicCell;
 use mio::event::Source;
 use mio::unix::SourceFd;
@@ -30,7 +30,7 @@ pub(crate) fn start_simple(
     #[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
     client_cipher: Cipher,
     server_cipher: Cipher,
-    up_counter: &mut SingleU64Adder,
+    up_counter: &U64Adder,
     device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
     compressor: Compressor,
     device_stop: DeviceStop,
@@ -85,7 +85,7 @@ fn start_simple0(
     #[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
     client_cipher: Cipher,
     server_cipher: Cipher,
-    up_counter: &mut SingleU64Adder,
+    up_counter: &U64Adder,
     device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
     compressor: Compressor,
 ) -> anyhow::Result<()> {

vnt/src/handle/tun_tap/windows.rs

@@ -7,7 +7,7 @@ use crate::handle::tun_tap::DeviceStop;
 use crate::handle::{CurrentDeviceInfo, PeerDeviceInfo};
 #[cfg(feature = "ip_proxy")]
 use crate::ip_proxy::IpProxyMap;
-use crate::util::{SingleU64Adder, StopManager};
+use crate::util::{StopManager, U64Adder};
 use crossbeam_utils::atomic::AtomicCell;
 use parking_lot::Mutex;
 use std::sync::Arc;
@@ -23,7 +23,7 @@ pub(crate) fn start_simple(
     #[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
     client_cipher: Cipher,
     server_cipher: Cipher,
-    up_counter: &mut SingleU64Adder,
+    up_counter: &U64Adder,
     device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
     compressor: Compressor,
     device_stop: DeviceStop,
@@ -76,7 +76,7 @@ fn start_simple0(
     #[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
     client_cipher: Cipher,
     server_cipher: Cipher,
-    up_counter: &mut SingleU64Adder,
+    up_counter: &U64Adder,
     device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
     compressor: Compressor,
 ) -> anyhow::Result<()> {

vnt/src/tun_tap_device/tun_create_helper.rs

@@ -16,7 +16,7 @@ use crate::handle::{CurrentDeviceInfo, PeerDeviceInfo};
 #[cfg(feature = "ip_proxy")]
 use crate::ip_proxy::IpProxyMap;
 use crate::tun_tap_device::vnt_device::DeviceWrite;
-use crate::util::{SingleU64Adder, StopManager};
+use crate::util::{StopManager, U64Adder};
 
 #[repr(transparent)]
 #[derive(Clone, Default)]
@@ -67,7 +67,7 @@ struct TunDeviceHelperInner {
     ip_proxy_map: Option<IpProxyMap>,
     client_cipher: Cipher,
     server_cipher: Cipher,
-    up_counter: SingleU64Adder,
+    up_counter: U64Adder,
     device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
     compressor: Compressor,
 }
@@ -81,7 +81,7 @@ impl TunDeviceHelper {
         #[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
         client_cipher: Cipher,
         server_cipher: Cipher,
-        up_counter: SingleU64Adder,
+        up_counter: U64Adder,
         device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
         compressor: Compressor,
         device_adapter: DeviceAdapter,

vnt/src/util/counter/adder.rs

@@ -1,139 +1,33 @@
-use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 
-/// 不安全的并发计数器，谨慎使用
+use crossbeam_utils::atomic::AtomicCell;
 
+#[derive(Clone, Default)]
 pub struct U64Adder {
-    global_index: Arc<AtomicUsize>,
+    count: Arc<AtomicCell<u64>>,
-    inner: Arc<U64AdderInner>,
-    index: usize,
-}
-#[derive(Clone)]
-pub struct SingleU64Adder {
-    inner: Arc<SingleU64AdderInner>,
-}
-impl SingleU64Adder {
-    pub fn new() -> Self {
-        Self {
-            inner: Arc::new(SingleU64AdderInner::new()),
-        }
-    }
-    pub fn add(&mut self, num: u64) {
-        self.inner.add(num);
-    }
-    pub fn get(&self) -> u64 {
-        self.inner.get()
-    }
-    pub fn watch(&self) -> WatchSingleU64Adder {
-        WatchSingleU64Adder {
-            inner: self.inner.clone(),
-        }
-    }
-}
-
-struct SingleU64AdderInner {
-    ptr: *mut u64,
-}
-
-impl SingleU64AdderInner {
-    fn new() -> Self {
-        Self {
-            ptr: Box::into_raw(Box::new(0)),
-        }
-    }
-    #[inline(always)]
-    fn add(&self, num: u64) {
-        unsafe { *self.ptr += num }
-    }
-
-    fn get(&self) -> u64 {
-        unsafe { *self.ptr }
-    }
-}
-impl Drop for SingleU64AdderInner {
-    fn drop(&mut self) {
-        unsafe {
-            let _ = Box::from_raw(self.ptr);
-        }
-    }
-}
-
-unsafe impl Send for SingleU64AdderInner {}
-
-unsafe impl Sync for SingleU64AdderInner {}
-
-struct U64AdderInner {
-    base: Vec<SingleU64AdderInner>,
-}
-
-impl U64AdderInner {
-    pub fn get(&self) -> u64 {
-        let mut count = 0;
-        for counter in self.base.iter() {
-            count += counter.get()
-        }
-        count
-    }
 }
 
 impl U64Adder {
-    /// 计数槽容量
-    pub fn with_capacity(capacity: usize) -> Self {
-        let mut base = Vec::with_capacity(capacity);
-        for _ in 0..capacity {
-            base.push(SingleU64AdderInner::new())
-        }
-        let inner = Arc::new(U64AdderInner { base });
-        U64Adder {
-            global_index: Arc::new(AtomicUsize::new(1)),
-            inner,
-            index: 0,
-        }
-    }
     pub fn add(&self, num: u64) {
-        self.inner.base[self.index].add(num);
+        self.count.fetch_add(num);
     }
     pub fn get(&self) -> u64 {
-        self.inner.get()
+        self.count.load()
     }
     pub fn watch(&self) -> WatchU64Adder {
         WatchU64Adder {
-            inner: self.inner.clone(),
+            count: self.count.clone(),
-        }
-    }
-}
-
-impl Clone for U64Adder {
-    fn clone(&self) -> Self {
-        let index = self.global_index.fetch_add(1, Ordering::AcqRel);
-        if index > self.inner.base.len() {
-            panic!()
-        }
-
-        Self {
-            global_index: self.global_index.clone(),
-            inner: self.inner.clone(),
-            index,
         }
     }
 }
 
 #[derive(Clone)]
 pub struct WatchU64Adder {
-    inner: Arc<U64AdderInner>,
+    count: Arc<AtomicCell<u64>>,
 }
 
 impl WatchU64Adder {
     pub fn get(&self) -> u64 {
-        self.inner.get()
+        self.count.load()
-    }
-}
-#[derive(Clone)]
-pub struct WatchSingleU64Adder {
-    inner: Arc<SingleU64AdderInner>,
-}
-impl WatchSingleU64Adder {
-    pub fn get(&self) -> u64 {
-        self.inner.get()
     }
 }