Let TCP peek methods take &self

This changes the `TcpStream::peek` and the various other peek methods to
take `&self` instead of `&mut self`, ensuring tokio API compatibility
and making them usable in more contexts.

To achieve this, we wrap the `ReadHalf` internals in `Mutex` and
`AtomicBool`, respectively. Note that we can't use `RefCell`/`Cell` here
because we also require `ReadHalf` to be `Sync`.

Author: teskje

src/net/tcp/split_owned.rs

@@ -40,20 +40,16 @@ impl OwnedReadHalf {
     /// Attempts to receive data on the socket, without removing that data from
     /// the queue, registering the current task for wakeup if data is not yet
     /// available.
-    pub fn poll_peek(
-        mut self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        buf: &mut ReadBuf,
-    ) -> Poll<io::Result<usize>> {
-        Pin::new(&mut self.inner).poll_peek(cx, buf)
+    pub fn poll_peek(&self, cx: &mut Context<'_>, buf: &mut ReadBuf) -> Poll<io::Result<usize>> {
+        self.inner.poll_peek(cx, buf)
     }
 
     /// Receives data on the socket from the remote address to which it is
     /// connected, without removing that data from the queue. On success,
     /// returns the number of bytes peeked.
     ///
     /// Successive calls return the same data.
-    pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+    pub async fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
         self.inner.peek(buf).await
     }
 }

src/net/tcp/stream.rs

@@ -5,7 +5,10 @@ use std::{
     io::{self, Error, Result},
     net::SocketAddr,
     pin::Pin,
-    sync::Arc,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc, Mutex,
+    },
     task::{ready, Context, Poll},
 };
 use tokio::{
@@ -40,11 +43,11 @@ impl TcpStream {
         let pair = Arc::new(pair);
         let read_half = ReadHalf {
             pair: pair.clone(),
-            rx: Rx {
+            rx: Mutex::new(Rx {
                 recv: receiver,
                 buffer: None,
-            },
-            is_closed: false,
+            }),
+            is_closed: AtomicBool::new(false),
         };
 
         let write_half = WriteHalf {
@@ -171,23 +174,23 @@ impl TcpStream {
     /// returns the number of bytes peeked.
     ///
     /// Successive calls return the same data.
-    pub async fn peek(&mut self, buf: &mut [u8]) -> Result<usize> {
+    pub async fn peek(&self, buf: &mut [u8]) -> Result<usize> {
         self.read_half.peek(buf).await
     }
 
     /// Attempts to receive data on the socket, without removing that data from
     /// the queue, registering the current task for wakeup if data is not yet
     /// available.
-    pub fn poll_peek(&mut self, cx: &mut Context<'_>, buf: &mut ReadBuf) -> Poll<Result<usize>> {
+    pub fn poll_peek(&self, cx: &mut Context<'_>, buf: &mut ReadBuf) -> Poll<Result<usize>> {
         self.read_half.poll_peek(cx, buf)
     }
 }
 
 pub(crate) struct ReadHalf {
     pub(crate) pair: Arc<SocketPair>,
-    rx: Rx,
+    rx: Mutex<Rx>,
     /// FIN received, EOF for reads
-    is_closed: bool,
+    is_closed: AtomicBool,
 }
 
 struct Rx {
@@ -200,27 +203,33 @@ struct Rx {
 }
 
 impl ReadHalf {
-    fn poll_read_priv(&mut self, cx: &mut Context<'_>, buf: &mut ReadBuf) -> Poll<Result<()>> {
-        if self.is_closed || buf.capacity() == 0 {
+    fn is_closed(&self) -> bool {
+        self.is_closed.load(Ordering::Acquire)
+    }
+
+    fn poll_read_priv(&self, cx: &mut Context<'_>, buf: &mut ReadBuf) -> Poll<Result<()>> {
+        if self.is_closed() || buf.capacity() == 0 {
             return Poll::Ready(Ok(()));
         }
 
-        if let Some(bytes) = self.rx.buffer.take() {
-            self.rx.buffer = Self::put_slice(bytes, buf);
+        let mut rx = self.rx.lock().unwrap();
+
+        if let Some(bytes) = rx.buffer.take() {
+            rx.buffer = Self::put_slice(bytes, buf);
 
             return Poll::Ready(Ok(()));
         }
 
-        match ready!(self.rx.recv.poll_recv(cx)) {
+        match ready!(rx.recv.poll_recv(cx)) {
             Some(seg) => {
                 tracing::trace!(target: TRACING_TARGET, src = ?self.pair.remote, dst = ?self.pair.local, protocol = %seg, "Recv");
 
                 match seg {
                     SequencedSegment::Data(bytes) => {
-                        self.rx.buffer = Self::put_slice(bytes, buf);
+                        rx.buffer = Self::put_slice(bytes, buf);
                     }
                     SequencedSegment::Fin => {
-                        self.is_closed = true;
+                        self.is_closed.store(true, Ordering::Release);
                     }
                 }
 
@@ -251,36 +260,34 @@ impl ReadHalf {
         }
     }
 
-    pub(crate) fn poll_peek(
-        &mut self,
-        cx: &mut Context<'_>,
-        buf: &mut ReadBuf,
-    ) -> Poll<Result<usize>> {
-        if self.is_closed || buf.capacity() == 0 {
+    pub(crate) fn poll_peek(&self, cx: &mut Context<'_>, buf: &mut ReadBuf) -> Poll<Result<usize>> {
+        if self.is_closed() || buf.capacity() == 0 {
             return Poll::Ready(Ok(0));
         }
 
+        let mut rx = self.rx.lock().unwrap();
+
         // If we have buffered data, peek from it
-        if let Some(bytes) = &self.rx.buffer {
+        if let Some(bytes) = &rx.buffer {
             let len = std::cmp::min(bytes.len(), buf.remaining());
             buf.put_slice(&bytes[..len]);
             return Poll::Ready(Ok(len));
         }
 
-        match ready!(self.rx.recv.poll_recv(cx)) {
+        match ready!(rx.recv.poll_recv(cx)) {
             Some(seg) => {
                 tracing::trace!(target: TRACING_TARGET, src = ?self.pair.remote, dst = ?self.pair.local, protocol = %seg, "Peek");
 
                 match seg {
                     SequencedSegment::Data(bytes) => {
                         let len = std::cmp::min(bytes.len(), buf.remaining());
                         buf.put_slice(&bytes[..len]);
-                        self.rx.buffer = Some(bytes);
+                        rx.buffer = Some(bytes);
 
                         Poll::Ready(Ok(len))
                     }
                     SequencedSegment::Fin => {
-                        self.is_closed = true;
+                        self.is_closed.store(true, Ordering::Release);
                         Poll::Ready(Ok(0))
                     }
                 }
@@ -292,7 +299,7 @@ impl ReadHalf {
         }
     }
 
-    pub(crate) async fn peek(&mut self, buf: &mut [u8]) -> Result<usize> {
+    pub(crate) async fn peek(&self, buf: &mut [u8]) -> Result<usize> {
         let mut buf = ReadBuf::new(buf);
         poll_fn(|cx| self.poll_peek(cx, &mut buf)).await
     }
@@ -302,7 +309,7 @@ impl Debug for ReadHalf {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("ReadHalf")
             .field("pair", &self.pair)
-            .field("is_closed", &self.is_closed)
+            .field("is_closed", &self.is_closed())
             .finish()
     }
 }
@@ -422,7 +429,7 @@ impl Debug for WriteHalf {
 
 impl AsyncRead for ReadHalf {
     fn poll_read(
-        mut self: Pin<&mut Self>,
+        self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         buf: &mut ReadBuf,
     ) -> Poll<Result<()>> {

tests/tcp.rs

@@ -821,7 +821,7 @@ fn peek_empty_buffer() -> Result {
     });
 
     sim.client("client", async move {
-        let mut s = TcpStream::connect(("server", PORT)).await?;
+        let s = TcpStream::connect(("server", PORT)).await?;
 
         // no-op peek with empty buffer
         let mut buf = [0; 0];