again

src/main.rs

@@ -7,7 +7,11 @@ use std::error::Error;
 use std::intrinsics::transmute;
 use std::io::{Read, Write};
 use std::mem::MaybeUninit;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
+use crossbeam_utils::thread;
+use grouping_by::GroupingBy;
+use pnet::packet::ipv4::Ipv4Packet;
+use socket2::Socket;
 
 #[repr(C)]
 pub struct Meta {
@@ -36,10 +40,6 @@ pub struct Config {
     pub local_secret: String,
     pub routers: Vec<ConfigRouter>,
 }
-use crossbeam_utils::thread;
-use grouping_by::GroupingBy;
-use pnet::packet::ipv4::Ipv4Packet;
-use socket2::Socket;
 
 fn main() -> Result<(), Box<dyn Error>> {
     let config: Config = serde_json::from_str(env::args().nth(1).ok_or("need param")?.as_str())?;
@@ -73,16 +73,42 @@ fn main() -> Result<(), Box<dyn Error>> {
         .collect();
     println!("created tuns");
 
+    // 获取系统的核心数量
+    let num_threads = std::thread::available_parallelism()
+        .map_or(1, |n| n.get()); // 默认使用1个线程，如果获取不到核心数
+
     thread::scope(|s| {
-        for router in router_readers.values_mut() {
-            s.spawn(|_| {
+        // 根据核心数量调整线程数
+        let readers_per_thread = (router_readers.len() as f32 / num_threads as f32).ceil() as usize;
+        let writers_per_thread = (router_writers3.len() as f32 / num_threads as f32).ceil() as usize;
+
+        // 将 router_readers 划分到多个线程
+        let readers_chunks: Vec<_> = router_readers
+            .into_iter()
+            .collect::<Vec<_>>()
+            .chunks(readers_per_thread)
+            .map(|chunk| chunk.to_vec())
+            .collect();
+
+        // 将 router_writers3 划分到多个线程
+        let writers_chunks: Vec<_> = router_writers3
+            .into_iter()
+            .collect::<Vec<_>>()
+            .chunks(writers_per_thread)
+            .map(|chunk| chunk.to_vec())
+            .collect();
+
+        // 启动处理 router_readers 的线程
+        for chunk in readers_chunks {
+            s.spawn(move |_| {
+                for router in chunk {
                     let mut buffer = [0u8; 1500 - 20]; // minus typical IP header space
-                let meta_size = size_of::<Meta>();
+                    let meta_size = std::mem::size_of::<Meta>();
 
                     // Pre-initialize with our Meta header (local -> remote)
                     let meta = Meta {
                         src_id: config.local_id,
-                    dst_id: router.config.remote_id,
+                        dst_id: router.1.config.remote_id,
                         reversed: 0,
                     };
                     // Turn the Meta struct into bytes
@@ -92,19 +118,22 @@ fn main() -> Result<(), Box<dyn Error>> {
                     buffer[..meta_size].copy_from_slice(meta_bytes);
 
                     loop {
-                    let n = router.tun_reader.read(&mut buffer[meta_size..]).unwrap();
-                    if let Some(ref addr) = *router.endpoint.read().unwrap() {
-                        router.encrypt(&mut buffer[meta_size..meta_size + n]);
+                        let n = router.1.tun_reader.read(&mut buffer[meta_size..]).unwrap();
+                        if let Some(ref addr) = *router.1.endpoint.read().unwrap() {
+                            router.1.encrypt(&mut buffer[meta_size..meta_size + n]);
                             #[cfg(target_os = "linux")]
-                        let _ = router.socket.set_mark(router.config.mark);
-                        let _ = router.socket.send_to(&buffer[..meta_size + n], addr);
+                            let _ = router.1.socket.set_mark(router.1.config.mark);
+                            let _ = router.1.socket.send_to(&buffer[..meta_size + n], addr);
+                        }
                     }
                 }
             });
         }
 
-        for (socket, mut router_writers) in router_writers3 {
+        // 启动处理 router_writers 的线程
+        for chunk in writers_chunks {
             s.spawn(move |_| {
+                for (socket, mut router_writers) in chunk {
                     let mut recv_buf = [MaybeUninit::uninit(); 1500];
                     loop {
                         let _ = (|| {
@@ -117,7 +146,7 @@ fn main() -> Result<(), Box<dyn Error>> {
                                 .split_at_mut_checked(header_len)
                                 .ok_or("malformed packet")?;
                             let (meta_bytes, payload) = rest
-                            .split_at_mut_checked(size_of::<Meta>())
+                                .split_at_mut_checked(std::mem::size_of::<Meta>())
                                 .ok_or("malformed packet")?;
                             let meta: &Meta = unsafe { transmute(meta_bytes.as_ptr()) };
                             if meta.dst_id == config.local_id && meta.reversed == 0 {
@@ -132,6 +161,7 @@ fn main() -> Result<(), Box<dyn Error>> {
                             Ok::<(), Box<dyn Error>>(())
                         })();
                     }
+                }
             });
         }
     })