fix(quic): adapt QuicServer::accpet more usages

Author: eareimu

gm-quic/src/server.rs

@@ -7,14 +7,10 @@ use std::{
 use dashmap::DashMap;
 use qbase::{
     cid::ConnectionId,
-    packet::{
-        header::{GetDcid, GetScid},
-        long, DataHeader, DataPacket, InitialHeader, RetryHeader,
-    },
+    packet::{header::GetScid, long, DataHeader, DataPacket, InitialHeader, RetryHeader},
     param::{Parameters, ServerParameters},
     sid::{handy::ConsistentConcurrency, ControlConcurrency},
     token::{ArcTokenRegistry, TokenProvider},
-    util::ArcAsyncDeque,
 };
 use qconnection::{conn::ArcConnection, path::Pathway, router::Router, usc::ArcUsc};
 use rustls::{
@@ -26,7 +22,7 @@ use rustls::{
 use crate::{get_or_create_usc, util, ConnKey, QuicConnection, CONNECTIONS};
 
 type TlsServerConfigBuilder<T> = ConfigBuilder<TlsServerConfig, T>;
-type QuicListner = ArcAsyncDeque<(Arc<QuicConnection>, Pathway)>;
+type QuicListner = Arc<util::Channel<(Arc<QuicConnection>, Pathway)>>;
 
 /// 理应全局只有一个server
 static SERVER: LazyLock<RwLock<Weak<QuicServer>>> = LazyLock::new(RwLock::default);
@@ -132,20 +128,18 @@ impl QuicServer {
     ///
     /// If all listening udp sockets are closed, this method will return an error.
     pub async fn accept(&self) -> io::Result<(Arc<QuicConnection>, Pathway)> {
-        // TODO: 错误
-        let error = || {
-            // let msg = "All listening udp sockets are closed";
-            // io::Error::new(io::ErrorKind::AddrNotAvailable, msg)
-            unreachable!()
+        let no_address_listening = || {
+            let error = "all listening udp sockets are closed";
+            io::Error::new(io::ErrorKind::AddrNotAvailable, error)
         };
-        // 这个VecDeque永远不会被close，所以这里不会是None
-        self.listener.pop().await.ok_or_else(error)
+        self.listener.recv().await.ok_or_else(no_address_listening)
     }
 }
 
 // internal methods
 impl QuicServer {
     pub(crate) fn try_to_accept_conn_from(mut packet: DataPacket, pathway: Pathway, usc: &ArcUsc) {
+        log::info!("try to accept connection from {}", pathway.dst_addr());
         let Some(server) = SERVER.read().unwrap().upgrade() else {
             return;
         };
@@ -161,13 +155,11 @@ impl QuicServer {
                 .unwrap();
         let (initial_dcid, client_initial_dcid) = match &mut packet.header {
             DataHeader::Long(long::DataHeader::Initial(hdr)) => {
-                let client_dcid = *hdr.get_dcid();
-                hdr.dcid = initial_scid;
+                let client_dcid = core::mem::replace(&mut hdr.dcid, initial_scid);
                 (*hdr.get_scid(), client_dcid)
             }
             DataHeader::Long(long::DataHeader::ZeroRtt(hdr)) => {
-                let client_dcid = *hdr.get_dcid();
-                hdr.dcid = initial_scid;
+                let client_dcid = core::mem::replace(&mut hdr.dcid, initial_scid);
                 (*hdr.get_scid(), client_dcid)
             }
             _ => return,
@@ -195,25 +187,25 @@ impl QuicServer {
             token_registry,
         );
         inner.add_initial_path(pathway, usc.clone());
-        let conn = QuicConnection {
+        let conn = Arc::new(QuicConnection {
             key: ConnKey::Server(initial_scid),
             inner: inner.clone(), // emm...
-        };
-        log::info!("incoming connection established");
-        server.listener.push_back((Arc::new(conn), pathway.filp()));
-        CONNECTIONS.insert(ConnKey::Server(initial_scid), inner);
-        _ = Router::try_to_route_packet_from(packet, pathway, usc);
+        });
+
+        if server.listener.send((conn, pathway.filp())).is_ok() {
+            CONNECTIONS.insert(ConnKey::Server(initial_scid), inner);
+            _ = Router::try_to_route_packet_from(packet, pathway, usc);
+        }
     }
 
     pub(crate) fn on_socket_close(addr: SocketAddr) {
         if let Some(server) = SERVER.read().unwrap().upgrade() {
             let bind_address_removed = server.sockets.remove(&addr).is_some();
-            // 所有已经绑定的地址都被关闭了，并且被动监听没有开启，那么就关闭server的监听...
-            // THINK: 不可能再接收新连接的情况下，Listener应该被立刻关闭？否
-            // if bind_address_removed && !server.passive_listening && server.sockets.is_empty() {
-            //     server.listener.close();
-            // }
-            _ = bind_address_removed;
+            // when: add listening sockets are removed, and passive listening is not enabled, it's not possiable
+            // to accept new connections anymore, so close the server's listener...
+            if bind_address_removed && !server.passive_listening && server.sockets.is_empty() {
+                server.listener.close();
+            }
         }
     }
 

gm-quic/src/util.rs

@@ -1,6 +1,7 @@
-use std::{io, path::Path};
+use std::{collections::VecDeque, io, path::Path, sync::Mutex};
 
 use rustls::pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer};
+use tokio::sync::Notify;
 
 pub fn parse_cert_files(
     cert_chain_file: impl AsRef<Path>,
@@ -22,3 +23,133 @@ pub fn parse_cert_files(
         .map_err(|e| cast_pem_error(e, key_file.as_ref()))?;
     Ok((cert_chain, key_der))
 }
+
+/// Unbound multi-producer multi-consumer (mpmc) channel.
+#[derive(Debug)]
+pub struct Channel<T> {
+    deque: Mutex<Option<VecDeque<T>>>,
+    notify: Notify,
+}
+
+impl<T> Channel<T> {
+    /// Create a new empty channel.
+    pub fn new() -> Self {
+        Self {
+            deque: Some(VecDeque::new()).into(),
+            notify: Default::default(),
+        }
+    }
+
+    /// Send an item to the channel.
+    ///
+    /// If the channel is closed, the item is returned as [`Err`].
+    pub fn send(&self, item: T) -> Result<(), T> {
+        log::info!("send {} to channel", core::any::type_name_of_val(&item));
+        let mut deque_guard = self.deque.lock().unwrap();
+
+        match deque_guard.as_mut() {
+            None => Err(item),
+            Some(deque) => {
+                deque.push_back(item);
+                drop(deque_guard);
+                self.notify.notify_one();
+                Ok(())
+            }
+        }
+    }
+
+    /// Close the channel.
+    ///
+    /// All unrecieved items will be dropped, and all pending [`Channel::recv`] calls will return [`None`].
+    pub fn close(&self) {
+        drop(self.deque.lock().unwrap().take());
+        self.notify.notify_waiters();
+    }
+
+    /// Try to recieve an item from the channel.
+    ///
+    /// If the channel is empty, [`TryRecvError::Empty`] is returned.
+    /// If the channel is closed, [`TryRecvError::Closed`] is returned.
+    pub fn try_recv(&self) -> Result<T, TryRecvError> {
+        let mut deque_guard = self.deque.lock().unwrap();
+        let deque = deque_guard.as_mut().ok_or(TryRecvError::Closed)?;
+        deque.pop_front().ok_or(TryRecvError::Empty)
+    }
+
+    /// Recieve an item from the channel.
+    ///
+    /// If the channel is empty, the call will wait until an item is sent.
+    ///
+    /// If the channel is closed, [`None`] is returned.
+    pub async fn recv(&self) -> Option<T> {
+        let fut = self.notify.notified();
+        tokio::pin!(fut);
+
+        loop {
+            // because the Notify::notify_one will only store one permit,
+            // we need to enable the notified to avoid missing the notification
+            fut.as_mut().enable();
+
+            match self.try_recv() {
+                Ok(item) => return Some(item),
+                Err(TryRecvError::Empty) => {
+                    fut.as_mut().await;
+                    fut.set(self.notify.notified());
+                }
+                Err(TryRecvError::Closed) => return None,
+            }
+        }
+    }
+}
+
+/// Error type for [`Channel::try_recv`].
+#[derive(thiserror::Error, Debug, Clone, Copy)]
+pub enum TryRecvError {
+    #[error("channel is empty")]
+    Empty,
+    #[error("channel is closed")]
+    Closed,
+}
+
+impl<T> Default for Channel<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn chennel() {
+        let ch = Arc::new(Channel::new());
+        let ch2 = ch.clone();
+        assert_eq!(ch.send(0), Ok(()));
+
+        let t1 = tokio::spawn(async move {
+            assert_eq!(ch.recv().await, Some(0));
+            assert_eq!(ch.recv().await, Some(1));
+            assert_eq!(ch.recv().await, Some(2));
+            assert_eq!(ch.recv().await, Some(3));
+            assert_eq!(ch.recv().await, None);
+        });
+
+        let t2 = tokio::spawn(async move {
+            assert_eq!(ch2.send(1), Ok(()));
+            tokio::task::yield_now().await;
+            assert_eq!(ch2.send(2), Ok(()));
+            tokio::task::yield_now().await;
+            assert_eq!(ch2.send(3), Ok(()));
+            tokio::task::yield_now().await;
+            ch2.close();
+            assert_eq!(ch2.send(4), Err(4));
+            assert_eq!(ch2.send(5), Err(5));
+        });
+
+        t1.await.unwrap();
+        t2.await.unwrap();
+    }
+}

qconnection/src/usc.rs

@@ -91,7 +91,7 @@ impl UscRegistry {
 /// This struct also provide useful methods to send datagrams via a given [`Pathway`].
 ///
 /// [`UdpSocketController`]: qudp::UdpSocketController
-#[derive(Clone, Deref)]
+#[derive(Debug, Clone, Deref)]
 pub struct ArcUsc {
     #[deref]
     usc: Arc<qudp::UdpSocketController>,