[PollStream/PollDataChannel] do not loose bytes (#341)

Returning Poll::Ready(Ok(n)) when write_fut is Some is incorrect since buf already contains new data and we should instead always return Poll::Ready(Ok(buf.len())) and update write_fut.

In the future, if/when we make write synchronous, this whole block could be simplified to just contain a single call (1 line) to write. See #344

Author: melekes

data/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 
 * Increased minimum support rust version to `1.60.0`.
+* Do not loose data in `PollDataChannel::poll_write` [#341](https://github.com/webrtc-rs/webrtc/pull/341).
 * `PollDataChannel::poll_shutdown`: make sure to flush any writes before shutting down [#340](https://github.com/webrtc-rs/webrtc/pull/340)
 
 ## 0.5.0

data/src/data_channel/mod.rs

@@ -566,44 +566,52 @@ impl AsyncWrite for PollDataChannel {
             return Poll::Ready(Ok(0));
         }
 
-        let (fut, fut_is_new) = match self.write_fut.as_mut() {
-            Some(fut) => (fut, false),
-            None => {
-                let data_channel = self.data_channel.clone();
-                let bytes = Bytes::copy_from_slice(buf);
-                (
-                    self.write_fut
-                        .get_or_insert(Box::pin(async move { data_channel.write(&bytes).await })),
-                    true,
-                )
+        if let Some(fut) = self.write_fut.as_mut() {
+            match fut.as_mut().poll(cx) {
+                Poll::Pending => Poll::Pending,
+                Poll::Ready(Err(e)) => {
+                    let data_channel = self.data_channel.clone();
+                    let bytes = Bytes::copy_from_slice(buf);
+                    self.write_fut =
+                        Some(Box::pin(async move { data_channel.write(&bytes).await }));
+                    Poll::Ready(Err(e.into()))
+                }
+                // Given the data is buffered, it's okay to ignore the number of written bytes.
+                //
+                // TODO: In the long term, `data_channel.write` should be made sync. Then we could
+                // remove the whole `if` condition and just call `data_channel.write`.
+                Poll::Ready(Ok(_)) => {
+                    let data_channel = self.data_channel.clone();
+                    let bytes = Bytes::copy_from_slice(buf);
+                    self.write_fut =
+                        Some(Box::pin(async move { data_channel.write(&bytes).await }));
+                    Poll::Ready(Ok(buf.len()))
+                }
             }
-        };
+        } else {
+            let data_channel = self.data_channel.clone();
+            let bytes = Bytes::copy_from_slice(buf);
+            let fut = self
+                .write_fut
+                .insert(Box::pin(async move { data_channel.write(&bytes).await }));
 
-        match fut.as_mut().poll(cx) {
-            Poll::Pending => {
+            match fut.as_mut().poll(cx) {
                 // If it's the first time we're polling the future, `Poll::Pending` can't be
-                // returned because that would mean the `PollStream` is not ready for writing. And
-                // this is not true since we've just created a future, which is going to write the
-                // buf to the underlying stream.
+                // returned because that would mean the `PollDataChannel` is not ready for writing.
+                // And this is not true since we've just created a future, which is going to write
+                // the buf to the underlying stream.
                 //
                 // It's okay to return `Poll::Ready` if the data is buffered (this is what the
                 // buffered writer and `File` do).
-                if fut_is_new {
-                    Poll::Ready(Ok(buf.len()))
-                } else {
-                    // If it's the subsequent poll, it's okay to return `Poll::Pending` as it
-                    // indicates that the `PollStream` is not ready for writing. Only one future
-                    // can be in progress at the time.
-                    Poll::Pending
+                Poll::Pending => Poll::Ready(Ok(buf.len())),
+                Poll::Ready(Err(e)) => {
+                    self.write_fut = None;
+                    Poll::Ready(Err(e.into()))
+                }
+                Poll::Ready(Ok(n)) => {
+                    self.write_fut = None;
+                    Poll::Ready(Ok(n))
                 }
-            }
-            Poll::Ready(Err(e)) => {
-                self.write_fut = None;
-                Poll::Ready(Err(e.into()))
-            }
-            Poll::Ready(Ok(n)) => {
-                self.write_fut = None;
-                Poll::Ready(Ok(n))
             }
         }
     }

sctp/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 
 * Increased minimum support rust version to `1.60.0`.
+* Do not loose data in `PollStream::poll_write` [#341](https://github.com/webrtc-rs/webrtc/pull/341).
 * `PollStream::poll_shutdown`: make sure to flush any writes before shutting down [#340](https://github.com/webrtc-rs/webrtc/pull/340)
 
 ## v0.6.1

sctp/src/stream/mod.rs

@@ -703,44 +703,54 @@ impl AsyncWrite for PollStream {
         cx: &mut Context<'_>,
         buf: &[u8],
     ) -> Poll<io::Result<usize>> {
-        let (fut, fut_is_new) = match self.write_fut.as_mut() {
-            Some(fut) => (fut, false),
-            None => {
-                let stream = self.stream.clone();
-                let bytes = Bytes::copy_from_slice(buf);
-                (
-                    self.write_fut
-                        .get_or_insert(Box::pin(async move { stream.write(&bytes).await })),
-                    true,
-                )
+        if buf.is_empty() {
+            return Poll::Ready(Ok(0));
+        }
+
+        if let Some(fut) = self.write_fut.as_mut() {
+            match fut.as_mut().poll(cx) {
+                Poll::Pending => Poll::Pending,
+                Poll::Ready(Err(e)) => {
+                    let stream = self.stream.clone();
+                    let bytes = Bytes::copy_from_slice(buf);
+                    self.write_fut = Some(Box::pin(async move { stream.write(&bytes).await }));
+                    Poll::Ready(Err(e.into()))
+                }
+                // Given the data is buffered, it's okay to ignore the number of written bytes.
+                //
+                // TODO: In the long term, `stream.write` should be made sync. Then we could
+                // remove the whole `if` condition and just call `stream.write`.
+                Poll::Ready(Ok(_)) => {
+                    let stream = self.stream.clone();
+                    let bytes = Bytes::copy_from_slice(buf);
+                    self.write_fut = Some(Box::pin(async move { stream.write(&bytes).await }));
+                    Poll::Ready(Ok(buf.len()))
+                }
             }
-        };
+        } else {
+            let stream = self.stream.clone();
+            let bytes = Bytes::copy_from_slice(buf);
+            let fut = self
+                .write_fut
+                .insert(Box::pin(async move { stream.write(&bytes).await }));
 
-        match fut.as_mut().poll(cx) {
-            Poll::Pending => {
+            match fut.as_mut().poll(cx) {
                 // If it's the first time we're polling the future, `Poll::Pending` can't be
                 // returned because that would mean the `PollStream` is not ready for writing. And
                 // this is not true since we've just created a future, which is going to write the
                 // buf to the underlying stream.
                 //
                 // It's okay to return `Poll::Ready` if the data is buffered (this is what the
                 // buffered writer and `File` do).
-                if fut_is_new {
-                    Poll::Ready(Ok(buf.len()))
-                } else {
-                    // If it's the subsequent poll, it's okay to return `Poll::Pending` as it
-                    // indicates that the `PollStream` is not ready for writing. Only one future
-                    // can be in progress at the time.
-                    Poll::Pending
+                Poll::Pending => Poll::Ready(Ok(buf.len())),
+                Poll::Ready(Err(e)) => {
+                    self.write_fut = None;
+                    Poll::Ready(Err(e.into()))
+                }
+                Poll::Ready(Ok(n)) => {
+                    self.write_fut = None;
+                    Poll::Ready(Ok(n))
                 }
-            }
-            Poll::Ready(Err(e)) => {
-                self.write_fut = None;
-                Poll::Ready(Err(e.into()))
-            }
-            Poll::Ready(Ok(n)) => {
-                self.write_fut = None;
-                Poll::Ready(Ok(n))
             }
         }
     }