perf: replace Mutex with RwLock

Author: iHsin

Cargo.toml

@@ -1,5 +1,6 @@
 [workspace]
 members = ["tuic", "tuic-quinn", "tuic-server", "tuic-client"]
+resolver = "2"
 
 [profile.release]
 opt-level = 3

tuic-client/src/connection/mod.rs

@@ -5,7 +5,6 @@ use crate::{
 };
 use crossbeam_utils::atomic::AtomicCell;
 use once_cell::sync::OnceCell;
-use parking_lot::Mutex;
 use quinn::{
     congestion::{BbrConfig, CubicConfig, NewRenoConfig},
     ClientConfig, Connection as QuinnConnection, Endpoint as QuinnEndpoint, EndpointConfig,
@@ -18,18 +17,17 @@ use std::{
     sync::{atomic::AtomicU32, Arc},
     time::Duration,
 };
-use tokio::{
-    sync::{Mutex as AsyncMutex, OnceCell as AsyncOnceCell},
-    time,
-};
+use tokio::sync::RwLock as AsyncRwLock;
+use tokio::{sync::OnceCell as AsyncOnceCell, time};
+
 use tuic_quinn::{side, Connection as Model};
 use uuid::Uuid;
 
 mod handle_stream;
 mod handle_task;
 
-static ENDPOINT: OnceCell<Mutex<Endpoint>> = OnceCell::new();
-static CONNECTION: AsyncOnceCell<AsyncMutex<Connection>> = AsyncOnceCell::const_new();
+static ENDPOINT: OnceCell<AsyncRwLock<Endpoint>> = OnceCell::new();
+static CONNECTION: AsyncOnceCell<AsyncRwLock<Connection>> = AsyncOnceCell::const_new();
 static TIMEOUT: AtomicCell<Duration> = AtomicCell::new(Duration::from_secs(0));
 
 pub const ERROR_CODE: VarInt = VarInt::from_u32(0);
@@ -117,7 +115,7 @@ impl Connection {
         };
 
         ENDPOINT
-            .set(Mutex::new(ep))
+            .set(AsyncRwLock::new(ep))
             .map_err(|_| "endpoint already initialized")
             .unwrap();
 
@@ -131,21 +129,22 @@ impl Connection {
             ENDPOINT
                 .get()
                 .unwrap()
-                .lock()
+                .read()
+                .await
                 .connect()
                 .await
-                .map(AsyncMutex::new)
+                .map(AsyncRwLock::new)
         };
 
         let try_get_conn = async {
             let mut conn = CONNECTION
                 .get_or_try_init(|| try_init_conn)
                 .await?
-                .lock()
+                .write()
                 .await;
 
             if conn.is_closed() {
-                let new_conn = ENDPOINT.get().unwrap().lock().connect().await?;
+                let new_conn = ENDPOINT.get().unwrap().read().await.connect().await?;
                 *conn = new_conn;
             }
 
@@ -254,7 +253,7 @@ struct Endpoint {
 }
 
 impl Endpoint {
-    async fn connect(&mut self) -> Result<Connection, Error> {
+    async fn connect(&self) -> Result<Connection, Error> {
         let mut last_err = None;
 
         for addr in self.server.resolve().await? {

tuic-server/src/connection/handle_task.rs

@@ -119,18 +119,24 @@ impl Connection {
                 src_addr = addr,
             );
 
-            let session = match self.udp_sessions.lock().entry(assoc_id) {
-                Entry::Occupied(entry) => entry.get().clone(),
-                Entry::Vacant(entry) => {
-                    let session = UdpSession::new(
-                        self.clone(),
-                        assoc_id,
-                        self.udp_relay_ipv6,
-                        self.max_external_pkt_size,
-                    )?;
-                    entry.insert(session.clone());
-                    session
-                }
+            let guard = self.udp_sessions.read().await;
+            let session = guard.get(&assoc_id).map(|v| v.to_owned());
+            drop(guard);
+            let session = match session {
+                Some(v) => v,
+                None => match self.udp_sessions.write().await.entry(assoc_id) {
+                    Entry::Occupied(entry) => entry.get().clone(),
+                    Entry::Vacant(entry) => {
+                        let session = UdpSession::new(
+                            self.clone(),
+                            assoc_id,
+                            self.udp_relay_ipv6,
+                            self.max_external_pkt_size,
+                        )?;
+                        entry.insert(session.clone());
+                        session
+                    }
+                },
             };
 
             let Some(socket_addr) = resolve_dns(&addr).await?.next() else {
@@ -162,8 +168,8 @@ impl Connection {
             user = self.auth,
         );
 
-        if let Some(session) = self.udp_sessions.lock().remove(&assoc_id) {
-            session.close();
+        if let Some(session) = self.udp_sessions.write().await.remove(&assoc_id) {
+            session.close().await;
         }
     }
 

tuic-server/src/connection/mod.rs

@@ -1,14 +1,14 @@
 use self::{authenticated::Authenticated, udp_session::UdpSession};
 use crate::{error::Error, utils::UdpRelayMode};
 use crossbeam_utils::atomic::AtomicCell;
-use parking_lot::Mutex;
 use quinn::{Connecting, Connection as QuinnConnection, VarInt};
 use register_count::Counter;
 use std::{
     collections::HashMap,
     sync::{atomic::AtomicU32, Arc},
     time::Duration,
 };
+use tokio::sync::RwLock as AsyncRwLock;
 use tokio::time;
 use tuic_quinn::{side, Authenticate, Connection as Model};
 use uuid::Uuid;
@@ -29,7 +29,7 @@ pub struct Connection {
     udp_relay_ipv6: bool,
     auth: Authenticated,
     task_negotiation_timeout: Duration,
-    udp_sessions: Arc<Mutex<HashMap<u16, UdpSession>>>,
+    udp_sessions: Arc<AsyncRwLock<HashMap<u16, UdpSession>>>,
     udp_relay_mode: Arc<AtomicCell<Option<UdpRelayMode>>>,
     max_external_pkt_size: usize,
     remote_uni_stream_cnt: Counter,
@@ -147,7 +147,7 @@ impl Connection {
             udp_relay_ipv6,
             auth: Authenticated::new(),
             task_negotiation_timeout,
-            udp_sessions: Arc::new(Mutex::new(HashMap::new())),
+            udp_sessions: Arc::new(AsyncRwLock::new(HashMap::new())),
             udp_relay_mode: Arc::new(AtomicCell::new(None)),
             max_external_pkt_size,
             remote_uni_stream_cnt: Counter::new(),
@@ -157,15 +157,15 @@ impl Connection {
         }
     }
 
-    fn authenticate(&self, auth: &Authenticate) -> Result<(), Error> {
+    async fn authenticate(&self, auth: &Authenticate) -> Result<(), Error> {
         if self.auth.get().is_some() {
             Err(Error::DuplicatedAuth)
         } else if self
             .users
             .get(&auth.uuid())
             .map_or(false, |password| auth.validate(password))
         {
-            self.auth.set(auth.uuid());
+            self.auth.set(auth.uuid()).await;
             Ok(())
         } else {
             Err(Error::AuthFailed(auth.uuid()))

tuic-server/src/connection/udp_session.rs

@@ -1,13 +1,13 @@
 use super::Connection;
 use crate::error::Error;
 use bytes::Bytes;
-use parking_lot::Mutex;
 use socket2::{Domain, Protocol, SockAddr, Socket, Type};
 use std::{
     io::Error as IoError,
     net::{Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket as StdUdpSocket},
     sync::Arc,
 };
+use tokio::sync::RwLock as AsyncRwLock;
 use tokio::{
     net::UdpSocket,
     sync::oneshot::{self, Sender},
@@ -23,7 +23,7 @@ struct UdpSessionInner {
     socket_v4: UdpSocket,
     socket_v6: Option<UdpSocket>,
     max_pkt_size: usize,
-    close: Mutex<Option<Sender<()>>>,
+    close: AsyncRwLock<Option<Sender<()>>>,
 }
 
 impl UdpSession {
@@ -89,7 +89,7 @@ impl UdpSession {
             socket_v4,
             socket_v6,
             max_pkt_size,
-            close: Mutex::new(Some(tx)),
+            close: AsyncRwLock::new(Some(tx)),
         }));
 
         let session_listening = session.clone();
@@ -161,7 +161,7 @@ impl UdpSession {
         }
     }
 
-    pub fn close(&self) {
-        let _ = self.0.close.lock().take().unwrap().send(());
+    pub async fn close(&self) {
+        let _ = self.0.close.write().await.take().unwrap().send(());
     }
 }