io: use Mutex instead of spinlock (#6403)

Author: maminrayej

tokio/src/io/split.rs

@@ -6,13 +6,11 @@
 
 use crate::io::{AsyncRead, AsyncWrite, ReadBuf};
 
-use std::cell::UnsafeCell;
 use std::fmt;
 use std::io;
 use std::pin::Pin;
-use std::sync::atomic::AtomicBool;
-use std::sync::atomic::Ordering::{Acquire, Release};
 use std::sync::Arc;
+use std::sync::Mutex;
 use std::task::{Context, Poll};
 
 cfg_io_util! {
@@ -38,8 +36,7 @@ cfg_io_util! {
         let is_write_vectored = stream.is_write_vectored();
 
         let inner = Arc::new(Inner {
-            locked: AtomicBool::new(false),
-            stream: UnsafeCell::new(stream),
+            stream: Mutex::new(stream),
             is_write_vectored,
         });
 
@@ -54,13 +51,19 @@ cfg_io_util! {
 }
 
 struct Inner<T> {
-    locked: AtomicBool,
-    stream: UnsafeCell<T>,
+    stream: Mutex<T>,
     is_write_vectored: bool,
 }
 
-struct Guard<'a, T> {
-    inner: &'a Inner<T>,
+impl<T> Inner<T> {
+    fn with_lock<R>(&self, f: impl FnOnce(Pin<&mut T>) -> R) -> R {
+        let mut guard = self.stream.lock().unwrap();
+
+        // safety: we do not move the stream.
+        let stream = unsafe { Pin::new_unchecked(&mut *guard) };
+
+        f(stream)
+    }
 }
 
 impl<T> ReadHalf<T> {
@@ -90,7 +93,7 @@ impl<T> ReadHalf<T> {
                 .ok()
                 .expect("`Arc::try_unwrap` failed");
 
-            inner.stream.into_inner()
+            inner.stream.into_inner().unwrap()
         } else {
             panic!("Unrelated `split::Write` passed to `split::Read::unsplit`.")
         }
@@ -111,8 +114,7 @@ impl<T: AsyncRead> AsyncRead for ReadHalf<T> {
         cx: &mut Context<'_>,
         buf: &mut ReadBuf<'_>,
     ) -> Poll<io::Result<()>> {
-        let mut inner = ready!(self.inner.poll_lock(cx));
-        inner.stream_pin().poll_read(cx, buf)
+        self.inner.with_lock(|stream| stream.poll_read(cx, buf))
     }
 }
 
@@ -122,67 +124,31 @@ impl<T: AsyncWrite> AsyncWrite for WriteHalf<T> {
         cx: &mut Context<'_>,
         buf: &[u8],
     ) -> Poll<Result<usize, io::Error>> {
-        let mut inner = ready!(self.inner.poll_lock(cx));
-        inner.stream_pin().poll_write(cx, buf)
+        self.inner.with_lock(|stream| stream.poll_write(cx, buf))
     }
 
     fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
-        let mut inner = ready!(self.inner.poll_lock(cx));
-        inner.stream_pin().poll_flush(cx)
+        self.inner.with_lock(|stream| stream.poll_flush(cx))
     }
 
     fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
-        let mut inner = ready!(self.inner.poll_lock(cx));
-        inner.stream_pin().poll_shutdown(cx)
+        self.inner.with_lock(|stream| stream.poll_shutdown(cx))
     }
 
     fn poll_write_vectored(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         bufs: &[io::IoSlice<'_>],
     ) -> Poll<Result<usize, io::Error>> {
-        let mut inner = ready!(self.inner.poll_lock(cx));
-        inner.stream_pin().poll_write_vectored(cx, bufs)
+        self.inner
+            .with_lock(|stream| stream.poll_write_vectored(cx, bufs))
     }
 
     fn is_write_vectored(&self) -> bool {
         self.inner.is_write_vectored
     }
 }
 
-impl<T> Inner<T> {
-    fn poll_lock(&self, cx: &mut Context<'_>) -> Poll<Guard<'_, T>> {
-        if self
-            .locked
-            .compare_exchange(false, true, Acquire, Acquire)
-            .is_ok()
-        {
-            Poll::Ready(Guard { inner: self })
-        } else {
-            // Spin... but investigate a better strategy
-
-            std::thread::yield_now();
-            cx.waker().wake_by_ref();
-
-            Poll::Pending
-        }
-    }
-}
-
-impl<T> Guard<'_, T> {
-    fn stream_pin(&mut self) -> Pin<&mut T> {
-        // safety: the stream is pinned in `Arc` and the `Guard` ensures mutual
-        // exclusion.
-        unsafe { Pin::new_unchecked(&mut *self.inner.stream.get()) }
-    }
-}
-
-impl<T> Drop for Guard<'_, T> {
-    fn drop(&mut self) {
-        self.inner.locked.store(false, Release);
-    }
-}
-
 unsafe impl<T: Send> Send for ReadHalf<T> {}
 unsafe impl<T: Send> Send for WriteHalf<T> {}
 unsafe impl<T: Sync> Sync for ReadHalf<T> {}