Add Selector::send_or_else to get the message back if not sent

src/select.rs

@@ -110,25 +110,42 @@ impl<'a, T> Selector<'a, T> {
     /// Add a send operation to the selector that sends the provided value.
     ///
     /// Once added, the selector can be used to run the provided handler function on completion of this operation.
-    pub fn send<U, F: FnMut(Result<(), SendError<U>>) -> T + 'a>(
+    pub fn send<U, F: FnOnce(Result<(), SendError<U>>) -> T + 'a>(
+        self,
+        sender: &'a Sender<U>,
+        msg: U,
+        mapper: F,
+    ) -> Self {
+        self.send_or_else(sender, msg, mapper, |_| {})
+    }
+
+    /// Add a send operation to the selector that sends the provided value.
+    ///
+    /// Once added, the selector can be used to run the provided handler function on completion of this operation.
+    ///
+    /// If another operation completes first, call the alternative handler instead.
+    pub fn send_or_else<U, F: FnOnce(Result<(), SendError<U>>) -> T + 'a, O: FnOnce(U) + 'a>(
         mut self,
         sender: &'a Sender<U>,
         msg: U,
         mapper: F,
+        or_else: O,
     ) -> Self {
-        struct SendSelection<'a, T, F, U> {
+        struct SendSelection<'a, T, F, O, U> {
             sender: &'a Sender<U>,
             msg: Option<U>,
             token: Token,
             signalled: Arc<Spinlock<VecDeque<Token>>>,
             hook: Option<Arc<Hook<U, SelectSignal>>>,
-            mapper: F,
+            mapper: Option<F>,
+            or_else: Option<O>,
             phantom: PhantomData<T>,
         }
 
-        impl<'a, T, F, U> Selection<'a, T> for SendSelection<'a, T, F, U>
+        impl<'a, T, F, O, U> Selection<'a, T> for SendSelection<'a, T, F, O, U>
         where
-            F: FnMut(Result<(), SendError<U>>) -> T,
+            F: FnOnce(Result<(), SendError<U>>) -> T,
+            O: FnOnce(U),
         {
             fn init(&mut self) -> Option<T> {
                 let token = self.token;
@@ -155,7 +172,7 @@ impl<'a, T> Selector<'a, T> {
                 );
 
                 if self.hook.is_none() {
-                    Some((self.mapper)(match r {
+                    Some((self.mapper.take().unwrap())(match r {
                         Ok(()) => Ok(()),
                         Err(TrySendTimeoutError::Disconnected(msg)) => Err(SendError(msg)),
                         _ => unreachable!(),
@@ -180,10 +197,14 @@ impl<'a, T> Selector<'a, T> {
                     return None;
                 };
 
-                Some((self.mapper)(res))
+                Some((self.mapper.take().unwrap())(res))
             }
 
             fn deinit(&mut self) {
+                if let Some(msg) = self.msg.take() {
+                    (self.or_else.take().unwrap())(msg);
+                }
+
                 if let Some(hook) = self.hook.take() {
                     // Remove hook
                     let hook: Arc<Hook<U, dyn Signal>> = hook;
@@ -193,6 +214,10 @@ impl<'a, T> Selector<'a, T> {
                         .unwrap()
                         .1
                         .retain(|s| s.signal().as_ptr() != hook.signal().as_ptr());
+
+                    if let Some(msg) = hook.try_take() {
+                        (self.or_else.take().unwrap())(msg);
+                    }
                 }
             }
         }
@@ -204,7 +229,8 @@ impl<'a, T> Selector<'a, T> {
             token,
             signalled: self.signalled.clone(),
             hook: None,
-            mapper,
+            mapper: Some(mapper),
+            or_else: Some(or_else),
             phantom: Default::default(),
         }));
 
@@ -214,7 +240,7 @@ impl<'a, T> Selector<'a, T> {
     /// Add a receive operation to the selector.
     ///
     /// Once added, the selector can be used to run the provided handler function on completion of this operation.
-    pub fn recv<U, F: FnMut(Result<U, RecvError>) -> T + 'a>(
+    pub fn recv<U, F: FnOnce(Result<U, RecvError>) -> T + 'a>(
         mut self,
         receiver: &'a Receiver<U>,
         mapper: F,
@@ -224,14 +250,14 @@ impl<'a, T> Selector<'a, T> {
             token: Token,
             signalled: Arc<Spinlock<VecDeque<Token>>>,
             hook: Option<Arc<Hook<U, SelectSignal>>>,
-            mapper: F,
+            mapper: Option<F>,
             received: bool,
             phantom: PhantomData<T>,
         }
 
         impl<'a, T, F, U> Selection<'a, T> for RecvSelection<'a, T, F, U>
         where
-            F: FnMut(Result<U, RecvError>) -> T,
+            F: FnOnce(Result<U, RecvError>) -> T,
         {
             fn init(&mut self) -> Option<T> {
                 let token = self.token;
@@ -254,7 +280,7 @@ impl<'a, T> Selector<'a, T> {
                 );
 
                 if self.hook.is_none() {
-                    Some((self.mapper)(match r {
+                    Some((self.mapper.take().unwrap())(match r {
                         Ok(msg) => Ok(msg),
                         Err(TryRecvTimeoutError::Disconnected) => Err(RecvError::Disconnected),
                         _ => unreachable!(),
@@ -274,7 +300,7 @@ impl<'a, T> Selector<'a, T> {
                     return None;
                 };
 
-                Some((self.mapper)(res))
+                Some((self.mapper.take().unwrap())(res))
             }
 
             fn deinit(&mut self) {
@@ -306,7 +332,7 @@ impl<'a, T> Selector<'a, T> {
             token,
             signalled: self.signalled.clone(),
             hook: None,
-            mapper,
+            mapper: Some(mapper),
             received: false,
             phantom: Default::default(),
         }));

tests/basic.rs

@@ -333,15 +333,19 @@ fn select_general() {
     let (tx0, rx0) = bounded(1);
     let (tx1, rx1) = unbounded();
 
-    for (i, t) in vec![tx0.clone(), tx1].into_iter().enumerate() {
+    for (i, t) in vec![tx0.clone(), tx1.clone()].into_iter().enumerate() {
         std::thread::spawn(move || {
             std::thread::sleep(std::time::Duration::from_millis(250));
             let _ = t.send(Foo(i));
         });
     }
 
+    let token = Box::new(());
     let x = Selector::new()
-        .recv(&rx0, |x| x)
+        .recv(&rx0, |x| {
+            drop(token);
+            x
+        })
         .recv(&rx1, |x| x)
         .wait()
         .unwrap();
@@ -360,9 +364,47 @@ fn select_general() {
         assert_eq!(rx0.recv().unwrap(), Foo(43));
     });
 
-    Selector::new().send(&tx0, Foo(43), |x| x).wait().unwrap();
+    let token = Box::new(());
+    Selector::new()
+        .send(&tx0, Foo(43), |x| {
+            drop(token);
+            x
+        })
+        .wait()
+        .unwrap();
 
     t.join().unwrap();
+
+    let mut else0 = None;
+    let mut else1 = None;
+    let res = Selector::new()
+        .send_or_else(
+            &tx0,
+            Foo(44),
+            |x| {
+                x.unwrap_err();
+                0
+            },
+            |x| else0 = Some(x),
+        )
+        .send_or_else(
+            &tx1,
+            Foo(45),
+            |x| {
+                x.unwrap();
+                1
+            },
+            |x| else1 = Some(x),
+        )
+        .wait();
+
+    if res == 0 {
+        assert_eq!(else0, None);
+        assert_eq!(else1, Some(Foo(45)));
+    } else {
+        assert_eq!(else0, Some(Foo(44)));
+        assert_eq!(else1, None);
+    }
 }
 
 struct MessageWithoutDebug(u32);