Add one-shot handlers to Router API.

Signed-off-by: Josh Matthews <[email protected]>

Author: jdm

src/router.rs

@@ -77,8 +77,32 @@ impl RouterProxy {
 
     /// Add a new `(receiver, callback)` pair to the router, and send a wakeup message
     /// to the router.
-    pub fn add_typed_route<T>(&self, receiver: IpcReceiver<T>, mut callback: TypedRouterHandler<T>)
-    where
+    pub fn add_typed_route<T>(
+        &self,
+        receiver: IpcReceiver<T>,
+        mut callback: TypedRouterMultiHandler<T>,
+    ) where
+        T: Serialize + for<'de> Deserialize<'de> + 'static,
+    {
+        // Before passing the message on to the callback, turn it into the appropriate type
+        let modified_callback = move |msg: IpcMessage| {
+            let typed_message = msg.to::<T>();
+            callback(typed_message)
+        };
+
+        self.add_route(
+            receiver.to_opaque(),
+            RouterHandler::Multi(Box::new(modified_callback)),
+        );
+    }
+
+    /// Add a new `(receiver, callback)` pair to the router, and send a wakeup message
+    /// to the router.
+    pub fn add_typed_one_shot_route<T>(
+        &self,
+        receiver: IpcReceiver<T>,
+        callback: TypedRouterOneShotHandler<T>,
+    ) where
         T: Serialize + for<'de> Deserialize<'de> + 'static,
     {
         // Before passing the message on to the callback, turn it into the appropriate type
@@ -87,7 +111,10 @@ impl RouterProxy {
             callback(typed_message)
         };
 
-        self.add_route(receiver.to_opaque(), Box::new(modified_callback));
+        self.add_route(
+            receiver.to_opaque(),
+            RouterHandler::Once(Box::new(modified_callback)),
+        );
     }
 
     /// Send a shutdown message to the router containing a ACK sender,
@@ -218,7 +245,18 @@ impl Router {
                     },
                     // Event from one of our registered receivers, call callback.
                     IpcSelectionResult::MessageReceived(id, message) => {
-                        self.handlers.get_mut(&id).unwrap()(message)
+                        match self.handlers.get_mut(&id).unwrap() {
+                            RouterHandler::Once(_) => {},
+                            RouterHandler::Multi(handler) => {
+                                (handler)(message);
+                                continue;
+                            },
+                        };
+                        let RouterHandler::Once(handler) = self.handlers.remove(&id).unwrap()
+                        else {
+                            continue;
+                        };
+                        (handler)(message);
                     },
                     IpcSelectionResult::ChannelClosed(id) => {
                         let _ = self.handlers.remove(&id).unwrap();
@@ -238,7 +276,18 @@ enum RouterMsg {
 }
 
 /// Function to call when a new event is received from the corresponding receiver.
-pub type RouterHandler = Box<dyn FnMut(IpcMessage) + Send>;
+pub type RouterMultiHandler = Box<dyn FnMut(IpcMessage) + Send>;
+
+/// Function to call once when a new event is received from the corresponding receiver.
+pub type RouterOneShotHandler = Box<dyn FnOnce(IpcMessage) + Send>;
+
+enum RouterHandler {
+    Once(RouterOneShotHandler),
+    Multi(RouterMultiHandler),
+}
+
+/// Like [RouterMultiHandler] but includes the type that will be passed to the callback
+pub type TypedRouterMultiHandler<T> = Box<dyn FnMut(Result<T, bincode::Error>) + Send>;
 
-/// Like [RouterHandler] but includes the type that will be passed to the callback
-pub type TypedRouterHandler<T> = Box<dyn FnMut(Result<T, bincode::Error>) + Send>;
+/// Like [RouterOneShotHandler] but includes the type that will be passed to the callback
+pub type TypedRouterOneShotHandler<T> = Box<dyn FnOnce(Result<T, bincode::Error>) + Send>;

src/test.rs

@@ -356,6 +356,29 @@ fn router_routing_to_new_crossbeam_receiver() {
     assert_eq!(received_person, person);
 }
 
+#[test]
+fn router_once_handler() {
+    let person = ("Patrick Walton".to_owned(), 29);
+    let (tx, rx) = ipc::channel().unwrap();
+    let (tx2, rx2) = ipc::channel().unwrap();
+
+    let router = RouterProxy::new();
+    let mut once_tx2 = Some(tx2);
+    router.add_typed_one_shot_route(
+        rx,
+        Box::new(move |_msg| once_tx2.take().unwrap().send(()).unwrap()),
+    );
+
+    // Send one single event.
+    tx.send(person.clone()).unwrap();
+    // Wait for acknowledgement that the callback ran.
+    rx2.recv().unwrap();
+    // This send should succeed but no handler should run. If it does run,
+    // a panic will occur.
+    tx.send(person.clone()).unwrap();
+    assert!(rx2.recv().is_err());
+}
+
 #[test]
 fn router_multiplexing() {
     let person = ("Patrick Walton".to_owned(), 29);