refactor: UdpSocketController to customized UdpSocket

qbase/src/time.rs

@@ -152,6 +152,8 @@ impl IdleTimer {
     // Checks health of the path and
     // determines whether a heartbeat packet needs to be sent.
     pub fn health(&mut self) -> Result<Option<PingFrame>, TimeOut> {
+        // TODO: 考虑连接还没关闭，闲置路径的 defer_idle_timeout 不能生效
+        //       得统一考虑成连接级的 defer_idle_timeout
         if let Some(t) = self.last_effective_comm {
             let elapsed = t.elapsed();
             if elapsed > self.idle_config.defer_idle_timeout() {

qbase/src/util.rs

@@ -14,4 +14,4 @@ mod unique_id;
 pub use unique_id::{UniqueId, UniqueIdGenerator};
 
 mod wakers;
-pub use wakers::{WakerVec, Wakers};
+pub use wakers::{WakerGroup, Wakers};

qbase/src/util/wakers.rs

@@ -8,30 +8,30 @@ use std::{
 use smallvec::SmallVec;
 
 #[derive(Debug, Clone)]
-pub struct WakerVec<const N: usize = 4> {
+pub struct WakerGroup<const N: usize = 4> {
     wakers: SmallVec<[Waker; N]>,
 }
 
-impl<const N: usize> Default for WakerVec<N> {
+impl<const N: usize> Default for WakerGroup<N> {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl<const N: usize> WakerVec<N> {
+impl<const N: usize> WakerGroup<N> {
     pub const fn new() -> Self {
         Self {
             wakers: SmallVec::new_const(),
         }
     }
 
-    pub fn register(&mut self, waker: &Waker) {
+    pub fn add(&mut self, waker: &Waker) {
         if !self.wakers.iter().any(|w| w.will_wake(waker)) {
             self.wakers.push(waker.clone());
         }
     }
 
-    pub fn unregister(&mut self, waker: &Waker) {
+    pub fn remove(&mut self, waker: &Waker) {
         self.wakers
             .retain(|registered| !registered.will_wake(waker));
     }
@@ -43,15 +43,15 @@ impl<const N: usize> WakerVec<N> {
     }
 }
 
-impl<const N: usize> Drop for WakerVec<N> {
+impl<const N: usize> Drop for WakerGroup<N> {
     fn drop(&mut self) {
         self.wake_all();
     }
 }
 
 #[derive(Debug)]
 pub struct Wakers<const N: usize = 4> {
-    wakers: Mutex<WakerVec<N>>,
+    inner: Mutex<WakerGroup<N>>,
 }
 
 impl<const N: usize> Wake for Wakers<N> {
@@ -73,24 +73,30 @@ impl<const N: usize> Default for Wakers<N> {
 impl<const N: usize> Wakers<N> {
     pub const fn new() -> Self {
         Self {
-            wakers: Mutex::new(WakerVec::new()),
+            inner: Mutex::new(WakerGroup::new()),
         }
     }
 
-    fn lock(&self) -> MutexGuard<'_, WakerVec<N>> {
-        self.wakers.lock().expect("Wakers mutex poisoned")
+    fn lock_guard(&self) -> MutexGuard<'_, WakerGroup<N>> {
+        self.inner.lock().expect("Wakers mutex poisoned")
     }
 
-    pub fn register(&self, waker: &Waker) {
-        self.lock().register(waker)
+    pub fn add(&self, waker: &Waker) {
+        self.lock_guard().add(waker)
     }
 
-    pub fn unregister(&self, waker: &Waker) {
-        self.lock().unregister(waker)
+    pub fn remove(&self, waker: &Waker) {
+        self.lock_guard().remove(waker)
+    }
+
+    pub fn together_with(self: &Arc<Self>, waker: &Waker) -> Waker {
+        let mut guard = self.lock_guard();
+        guard.add(waker);
+        Waker::from(self.clone())
     }
 
     pub fn wake_all(&self) {
-        { mem::replace(&mut *self.lock(), WakerVec::new()) }.wake_all()
+        { mem::replace(&mut *self.lock_guard(), WakerGroup::new()) }.wake_all()
     }
 
     pub fn to_waker(self: &Arc<Self>) -> Waker {
@@ -102,10 +108,10 @@ impl<const N: usize> Wakers<N> {
         cx: &mut Context<'_>,
         poll: impl FnOnce(&mut Context<'_>) -> Poll<T>,
     ) -> Poll<T> {
-        self.register(cx.waker());
+        self.add(cx.waker());
         let result = poll(&mut Context::from_waker(&self.to_waker()));
         if result.is_ready() {
-            self.unregister(cx.waker());
+            self.remove(cx.waker());
         }
         result
     }

qinterface/src/io/handy.rs

@@ -84,7 +84,7 @@ pub mod qudp {
             }
         }
 
-        fn usc(&self) -> io::Result<&qudp::UdpSocket> {
+        fn socket(&self) -> io::Result<&qudp::UdpSocket> {
             self.io
                 .as_ref()
                 .map_err(|e| io::Error::from(e.clone()))
@@ -98,7 +98,7 @@ pub mod qudp {
         }
 
         fn bound_addr(&self) -> io::Result<SocketAddr> {
-            self.usc()?.local_addr()
+            self.socket()?.local_addr()
         }
 
         fn max_segments(&self) -> io::Result<usize> {
@@ -115,11 +115,17 @@ pub mod qudp {
             pkts: &[io::IoSlice],
             route: Route,
         ) -> Poll<io::Result<usize>> {
-            let io = self.usc()?;
-            self.send_wakers.combine_with(cx, |cx| {
-                debug_assert_eq!(route.ecn(), None);
-                io.poll_send(cx, pkts, &route.line)
-            })
+            let io = self.socket()?;
+            let waker = cx.waker();
+            let waker_group = self.send_wakers.together_with(waker);
+            let cx = &mut Context::from_waker(&waker_group);
+
+            debug_assert_eq!(route.ecn(), None);
+            let result = io.poll_send(cx, pkts, &route);
+            if result.is_ready() {
+                self.send_wakers.remove(waker);
+            }
+            result
         }
 
         fn poll_recv(
@@ -128,7 +134,7 @@ pub mod qudp {
             pkts: &mut [BytesMut],
             route: &mut [Route],
         ) -> Poll<io::Result<usize>> {
-            let io = self.usc()?;
+            let io = self.socket()?;
             self.recv_wakers.combine_with(cx, |cx| {
                 let len = route.len().min(pkts.len());
                 let mut rcvd_lines = Vec::with_capacity(len);
@@ -152,7 +158,7 @@ pub mod qudp {
         }
 
         fn poll_close(&mut self, _cx: &mut Context) -> Poll<io::Result<()>> {
-            self.usc()?;
+            self.socket()?;
             self.send_wakers.wake_all();
             self.recv_wakers.wake_all();
             self.io = Ok(Err(Closed(())));

qprotocol/Cargo.toml

@@ -12,9 +12,13 @@ categories.workspace = true
 [dependencies]
 bytes = { workspace = true }
 nom = { workspace = true }
+dashmap = { workspace = true }
 qbase = { workspace = true }
+qinterface = { workspace = true }
+qudp = { workspace = true }
 rand = { workspace = true }
 thiserror = { workspace = true }
+tokio = { workspace = true }
 
 [features]
 # Enable shorter TTL only for tests (especially integration tests in other crates).