fix(s2n-quic-dc): use wake_forced for worker::Waker (#2415)

Author: camshaft

dc/s2n-quic-dc-benches/src/streams.rs

@@ -37,7 +37,10 @@ fn pair(
     accept_flavor: accept::Flavor,
 ) -> (stream::testing::Client, stream::testing::Server) {
     let client = stream::testing::Client::default();
-    let server = stream::testing::Server::new(protocol, accept_flavor);
+    let server = stream::testing::Server::builder()
+        .protocol(protocol)
+        .accept_flavor(accept_flavor)
+        .build();
     client.handshake_with(&server).unwrap();
     (client, server)
 }

dc/s2n-quic-dc/src/stream.rs

@@ -28,6 +28,8 @@ pub mod socket;
 
 #[cfg(any(test, feature = "testing"))]
 pub mod testing;
+#[cfg(test)]
+mod tests;
 
 bitflags::bitflags! {
     #[derive(Clone, Copy, Debug, PartialEq, Eq)]

dc/s2n-quic-dc/src/stream/send/tests.rs

@@ -9,7 +9,7 @@ use tracing::Instrument as _;
 
 fn pair(protocol: Protocol) -> (testing::Client, testing::Server) {
     let client = testing::Client::default();
-    let server = testing::Server::new(protocol, Default::default());
+    let server = testing::Server::builder().protocol(protocol).build();
     (client, server)
 }
 

dc/s2n-quic-dc/src/stream/server/tokio/tcp.rs

@@ -79,7 +79,7 @@ where
         let mut context = worker::Context::new(&self);
 
         poll_fn(move |cx| {
-            workers.update_task_context(cx);
+            workers.poll_start(cx);
 
             let now = self.env.clock().get_time();
             let publisher = publisher(&self.subscriber, &now);

dc/s2n-quic-dc/src/stream/server/tokio/tcp/manager.rs

@@ -31,7 +31,6 @@ where
 {
     inner: Inner<W>,
     waker_set: waker::Set,
-    root_waker: Option<Waker>,
 }
 
 /// Split the tasks from the waker set to avoid ownership issues
@@ -120,11 +119,7 @@ where
             sojourn_time: RttEstimator::new(Duration::from_secs(30)),
         };
 
-        Self {
-            inner,
-            waker_set,
-            root_waker: None,
-        }
+        Self { inner, waker_set }
     }
 
     #[inline]
@@ -149,19 +144,8 @@ where
 
     /// Must be called before polling any workers
     #[inline]
-    pub fn update_task_context(&mut self, cx: &mut task::Context) {
-        let new_waker = cx.waker();
-
-        let root_task_requires_update = if let Some(waker) = self.root_waker.as_ref() {
-            !waker.will_wake(new_waker)
-        } else {
-            true
-        };
-
-        if root_task_requires_update {
-            self.waker_set.update_root(new_waker);
-            self.root_waker = Some(new_waker.clone());
-        }
+    pub fn poll_start(&mut self, cx: &mut task::Context) {
+        self.waker_set.poll_start(cx);
     }
 
     #[inline]
@@ -221,8 +205,15 @@ where
         Pub: EndpointPublisher,
         C: Clock,
     {
+        let ready = self.waker_set.drain();
+
+        // no need to actually poll any workers if none are active
+        if self.inner.by_sojourn_time.is_empty() {
+            return ControlFlow::Continue(());
+        }
+
         // poll any workers that are ready
-        for idx in self.waker_set.drain() {
+        for idx in ready {
             if self.inner.poll_worker(idx, cx, publisher, clock).is_break() {
                 return ControlFlow::Break(());
             }