adding branching wakers + fixing the arc construction

libp2p · Dec 22, 2023 · 59dfdf8 · 59dfdf8
1 parent d6b026c
commit 59dfdf8
Show file tree

Hide file tree

Showing 2 changed files with 250 additions and 208 deletions.
diff --git a/muxers/yamux/src/lib.rs b/muxers/yamux/src/lib.rs
@@ -129,25 +129,23 @@ where
     ) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
         let this = self.get_mut();
 
-        let inbound_stream = ready!(this.poll_inner(cx))?;
-
-        if this.inbound_stream_buffer.len() >= MAX_BUFFERED_INBOUND_STREAMS {
-            tracing::warn!(
-                stream=%inbound_stream.0,
-                "dropping stream because buffer is full"
-            );
-            drop(inbound_stream);
-        } else {
-            this.inbound_stream_buffer.push_back(inbound_stream);
-
-            if let Some(waker) = this.inbound_stream_waker.take() {
-                waker.wake()
+        loop {
+            let inbound_stream = ready!(this.poll_inner(cx))?;
+
+            if this.inbound_stream_buffer.len() >= MAX_BUFFERED_INBOUND_STREAMS {
+                tracing::warn!(
+                    stream=%inbound_stream.0,
+                    "dropping stream because buffer is full"
+                );
+                drop(inbound_stream);
+            } else {
+                this.inbound_stream_buffer.push_back(inbound_stream);
+
+                if let Some(waker) = this.inbound_stream_waker.take() {
+                    waker.wake()
+                }
             }
         }
-
-        // Schedule an immediate wake-up, allowing other code to run.
-        // cx.waker().wake_by_ref();
-        Poll::Pending
     }
 }