use blocking state in select macro

Author: ibraheemdev

crossbeam-channel/src/channel.rs

@@ -14,6 +14,7 @@ use crate::err::{
 };
 use crate::flavors;
 use crate::select::{Operation, SelectHandle, Token};
+use crate::waker::BlockingState;
 
 /// Creates a multi-producer multi-consumer channel of unbounded capacity.
 ///
@@ -1358,6 +1359,14 @@ impl<T> fmt::Debug for IntoIter<T> {
 }
 
 impl<T> SelectHandle for Sender<T> {
+    fn start(&self) -> Option<BlockingState<'_>> {
+        match &self.flavor {
+            SenderFlavor::Array(chan) => chan.sender().start_ref(),
+            SenderFlavor::List(chan) => chan.sender().start_ref(),
+            SenderFlavor::Zero(chan) => chan.start(),
+        }
+    }
+
     fn try_select(&self, token: &mut Token) -> bool {
         match &self.flavor {
             SenderFlavor::Array(chan) => chan.sender().try_select(token),
@@ -1370,11 +1379,11 @@ impl<T> SelectHandle for Sender<T> {
         None
     }
 
-    fn register(&self, oper: Operation, cx: &Context) -> bool {
+    fn register(&self, oper: Operation, cx: &Context, state: Option<&BlockingState<'_>>) -> bool {
         match &self.flavor {
-            SenderFlavor::Array(chan) => chan.sender().register(oper, cx),
-            SenderFlavor::List(chan) => chan.sender().register(oper, cx),
-            SenderFlavor::Zero(chan) => chan.sender().register(oper, cx),
+            SenderFlavor::Array(chan) => chan.sender().register(oper, cx, state),
+            SenderFlavor::List(chan) => chan.sender().register(oper, cx, state),
+            SenderFlavor::Zero(chan) => chan.sender().register(oper, cx, state),
         }
     }
 
@@ -1420,6 +1429,17 @@ impl<T> SelectHandle for Sender<T> {
 }
 
 impl<T> SelectHandle for Receiver<T> {
+    fn start(&self) -> Option<BlockingState<'_>> {
+        match &self.flavor {
+            ReceiverFlavor::Array(chan) => chan.receiver().start_ref(),
+            ReceiverFlavor::List(chan) => chan.receiver().start_ref(),
+            ReceiverFlavor::Zero(chan) => chan.start(),
+            ReceiverFlavor::At(chan) => chan.start(),
+            ReceiverFlavor::Tick(chan) => chan.start(),
+            ReceiverFlavor::Never(chan) => chan.start(),
+        }
+    }
+
     fn try_select(&self, token: &mut Token) -> bool {
         match &self.flavor {
             ReceiverFlavor::Array(chan) => chan.receiver().try_select(token),
@@ -1442,14 +1462,14 @@ impl<T> SelectHandle for Receiver<T> {
         }
     }
 
-    fn register(&self, oper: Operation, cx: &Context) -> bool {
+    fn register(&self, oper: Operation, cx: &Context, state: Option<&BlockingState<'_>>) -> bool {
         match &self.flavor {
-            ReceiverFlavor::Array(chan) => chan.receiver().register(oper, cx),
-            ReceiverFlavor::List(chan) => chan.receiver().register(oper, cx),
-            ReceiverFlavor::Zero(chan) => chan.receiver().register(oper, cx),
-            ReceiverFlavor::At(chan) => chan.register(oper, cx),
-            ReceiverFlavor::Tick(chan) => chan.register(oper, cx),
-            ReceiverFlavor::Never(chan) => chan.register(oper, cx),
+            ReceiverFlavor::Array(chan) => chan.receiver().register(oper, cx, state),
+            ReceiverFlavor::List(chan) => chan.receiver().register(oper, cx, state),
+            ReceiverFlavor::Zero(chan) => chan.receiver().register(oper, cx, state),
+            ReceiverFlavor::At(chan) => chan.register(oper, cx, state),
+            ReceiverFlavor::Tick(chan) => chan.register(oper, cx, state),
+            ReceiverFlavor::Never(chan) => chan.register(oper, cx, state),
         }
     }
 

crossbeam-channel/src/flavors/array.rs

@@ -20,7 +20,7 @@ use crossbeam_utils::{Backoff, CachePadded};
 use crate::context::Context;
 use crate::err::{RecvTimeoutError, SendTimeoutError, TryRecvError, TrySendError};
 use crate::select::{Operation, SelectHandle, Selected, Token};
-use crate::waker::SyncWaker;
+use crate::waker::{BlockingState, SyncWaker};
 
 /// A slot in a channel.
 struct Slot<T> {
@@ -401,7 +401,7 @@ impl<T> Channel<T> {
             Context::with(|cx| {
                 // Prepare for blocking until a receiver wakes us up.
                 let oper = Operation::hook(token);
-                self.senders.register2(oper, cx, &state);
+                self.senders.register(oper, cx, &state);
 
                 // Has the channel become ready just now?
                 if !self.is_full() || self.is_disconnected() {
@@ -478,7 +478,7 @@ impl<T> Channel<T> {
             Context::with(|cx| {
                 // Prepare for blocking until a sender wakes us up.
                 let oper = Operation::hook(token);
-                self.receivers.register2(oper, cx, &mut state);
+                self.receivers.register(oper, cx, &state);
 
                 // Has the channel become ready just now?
                 if !self.is_empty() || self.is_disconnected() {
@@ -629,7 +629,18 @@ pub(crate) struct Receiver<'a, T>(&'a Channel<T>);
 /// Sender handle to a channel.
 pub(crate) struct Sender<'a, T>(&'a Channel<T>);
 
+impl<'a, T> Receiver<'a, T> {
+    /// Same as `SelectHandle::start`, but with a more specific lifetime.
+    pub(crate) fn start_ref(&self) -> Option<BlockingState<'a>> {
+        Some(self.0.receivers.start())
+    }
+}
+
 impl<T> SelectHandle for Receiver<'_, T> {
+    fn start(&self) -> Option<BlockingState<'_>> {
+        self.start_ref()
+    }
+
     fn try_select(&self, token: &mut Token) -> bool {
         self.0.start_recv(token) == Status::Ready
     }
@@ -638,8 +649,9 @@ impl<T> SelectHandle for Receiver<'_, T> {
         None
     }
 
-    fn register(&self, oper: Operation, cx: &Context) -> bool {
-        self.0.receivers.register(oper, cx);
+    fn register(&self, oper: Operation, cx: &Context, state: Option<&BlockingState<'_>>) -> bool {
+        let state = state.expect("Receiver::start returns blocking state");
+        self.0.receivers.register(oper, cx, state);
         self.is_ready()
     }
 
@@ -665,7 +677,18 @@ impl<T> SelectHandle for Receiver<'_, T> {
     }
 }
 
+impl<'a, T> Sender<'a, T> {
+    /// Same as `SelectHandle::start`, but with a more specific lifetime.
+    pub(crate) fn start_ref(&self) -> Option<BlockingState<'a>> {
+        Some(self.0.senders.start())
+    }
+}
+
 impl<T> SelectHandle for Sender<'_, T> {
+    fn start(&self) -> Option<BlockingState<'_>> {
+        self.start_ref()
+    }
+
     fn try_select(&self, token: &mut Token) -> bool {
         self.0.start_send(token) == Status::Ready
     }
@@ -674,8 +697,9 @@ impl<T> SelectHandle for Sender<'_, T> {
         None
     }
 
-    fn register(&self, oper: Operation, cx: &Context) -> bool {
-        self.0.senders.register(oper, cx);
+    fn register(&self, oper: Operation, cx: &Context, state: Option<&BlockingState<'_>>) -> bool {
+        let state = state.expect("Sender::start returns blocking state");
+        self.0.senders.register(oper, cx, state);
         self.is_ready()
     }
 

crossbeam-channel/src/flavors/at.rs

@@ -10,6 +10,7 @@ use crate::context::Context;
 use crate::err::{RecvTimeoutError, TryRecvError};
 use crate::select::{Operation, SelectHandle, Token};
 use crate::utils;
+use crate::waker::BlockingState;
 
 /// Result of a receive operation.
 pub(crate) type AtToken = Option<Instant>;
@@ -140,6 +141,10 @@ impl Channel {
 }
 
 impl SelectHandle for Channel {
+    fn start(&self) -> Option<BlockingState<'_>> {
+        None
+    }
+
     #[inline]
     fn try_select(&self, token: &mut Token) -> bool {
         match self.try_recv() {
@@ -166,7 +171,12 @@ impl SelectHandle for Channel {
     }
 
     #[inline]
-    fn register(&self, _oper: Operation, _cx: &Context) -> bool {
+    fn register(
+        &self,
+        _oper: Operation,
+        _cx: &Context,
+        _state: Option<&BlockingState<'_>>,
+    ) -> bool {
         self.is_ready()
     }
 

crossbeam-channel/src/flavors/list.rs

@@ -13,7 +13,7 @@ use crossbeam_utils::{Backoff, CachePadded};
 use crate::context::Context;
 use crate::err::{RecvTimeoutError, SendTimeoutError, TryRecvError, TrySendError};
 use crate::select::{Operation, SelectHandle, Selected, Token};
-use crate::waker::SyncWaker;
+use crate::waker::{BlockingState, SyncWaker};
 
 // TODO(stjepang): Once we bump the minimum required Rust version to 1.28 or newer, re-apply the
 // following changes by @kleimkuhler:
@@ -510,7 +510,7 @@ impl<T> Channel<T> {
             // Prepare for blocking until a sender wakes us up.
             Context::with(|cx| {
                 let oper = Operation::hook(token);
-                self.receivers.register2(oper, cx, &state);
+                self.receivers.register(oper, cx, &state);
 
                 // Has the channel become ready just now?
                 if !self.is_empty() || self.is_disconnected() {
@@ -734,7 +734,18 @@ pub(crate) struct Receiver<'a, T>(&'a Channel<T>);
 /// Sender handle to a channel.
 pub(crate) struct Sender<'a, T>(&'a Channel<T>);
 
+impl<'a, T> Receiver<'a, T> {
+    /// Same as `SelectHandle::start`, but with a more specific lifetime.
+    pub(crate) fn start_ref(&self) -> Option<BlockingState<'a>> {
+        Some(self.0.receivers.start())
+    }
+}
+
 impl<T> SelectHandle for Receiver<'_, T> {
+    fn start(&self) -> Option<BlockingState<'_>> {
+        self.start_ref()
+    }
+
     fn try_select(&self, token: &mut Token) -> bool {
         self.0.start_recv(token) == Status::Ready
     }
@@ -743,8 +754,9 @@ impl<T> SelectHandle for Receiver<'_, T> {
         None
     }
 
-    fn register(&self, oper: Operation, cx: &Context) -> bool {
-        self.0.receivers.register(oper, cx);
+    fn register(&self, oper: Operation, cx: &Context, state: Option<&BlockingState<'_>>) -> bool {
+        let state = state.expect("Receiver::start returns blocking stat");
+        self.0.receivers.register(oper, cx, state);
         self.is_ready()
     }
 
@@ -770,7 +782,18 @@ impl<T> SelectHandle for Receiver<'_, T> {
     }
 }
 
-impl<T> SelectHandle for Sender<'_, T> {
+impl<'a, T> Sender<'a, T> {
+    /// Same as `SelectHandle::start`, but with a more specific lifetime.
+    pub(crate) fn start_ref(&self) -> Option<BlockingState<'a>> {
+        None
+    }
+}
+
+impl<'a, T> SelectHandle for Sender<'a, T> {
+    fn start(&self) -> Option<BlockingState<'a>> {
+        None
+    }
+
     fn try_select(&self, token: &mut Token) -> bool {
         self.0.start_send(token)
     }
@@ -779,7 +802,12 @@ impl<T> SelectHandle for Sender<'_, T> {
         None
     }
 
-    fn register(&self, _oper: Operation, _cx: &Context) -> bool {
+    fn register(
+        &self,
+        _oper: Operation,
+        _cx: &Context,
+        _state: Option<&BlockingState<'_>>,
+    ) -> bool {
         self.is_ready()
     }
 

crossbeam-channel/src/flavors/never.rs

@@ -9,6 +9,7 @@ use crate::context::Context;
 use crate::err::{RecvTimeoutError, TryRecvError};
 use crate::select::{Operation, SelectHandle, Token};
 use crate::utils;
+use crate::waker::BlockingState;
 
 /// This flavor doesn't need a token.
 pub(crate) type NeverToken = ();
@@ -72,6 +73,10 @@ impl<T> Channel<T> {
 }
 
 impl<T> SelectHandle for Channel<T> {
+    fn start(&self) -> Option<BlockingState<'_>> {
+        None
+    }
+
     #[inline]
     fn try_select(&self, _token: &mut Token) -> bool {
         false
@@ -83,7 +88,12 @@ impl<T> SelectHandle for Channel<T> {
     }
 
     #[inline]
-    fn register(&self, _oper: Operation, _cx: &Context) -> bool {
+    fn register(
+        &self,
+        _oper: Operation,
+        _cx: &Context,
+        _state: Option<&BlockingState<'_>>,
+    ) -> bool {
         self.is_ready()
     }
 

crossbeam-channel/src/flavors/tick.rs

@@ -10,6 +10,7 @@ use crossbeam_utils::atomic::AtomicCell;
 use crate::context::Context;
 use crate::err::{RecvTimeoutError, TryRecvError};
 use crate::select::{Operation, SelectHandle, Token};
+use crate::waker::BlockingState;
 
 /// Result of a receive operation.
 pub(crate) type TickToken = Option<Instant>;
@@ -115,6 +116,10 @@ impl Channel {
 }
 
 impl SelectHandle for Channel {
+    fn start(&self) -> Option<BlockingState<'_>> {
+        None
+    }
+
     #[inline]
     fn try_select(&self, token: &mut Token) -> bool {
         match self.try_recv() {
@@ -136,7 +141,12 @@ impl SelectHandle for Channel {
     }
 
     #[inline]
-    fn register(&self, _oper: Operation, _cx: &Context) -> bool {
+    fn register(
+        &self,
+        _oper: Operation,
+        _cx: &Context,
+        _state: Option<&BlockingState<'_>>,
+    ) -> bool {
         self.is_ready()
     }