From b93a22423bd230291086a97204537346dd0a537e Mon Sep 17 00:00:00 2001
From: wngr <oliver@wngr.de>
Date: Mon, 7 Dec 2020 17:08:31 +0100
Subject: [PATCH] feat: Expose `timeout` parameter to the main event loop, such
 that control flow is passed back to an external caller. This provides more
 flexibility to the API user.

Related #43 #38
---
 examples/systray-example.rs |  2 +-
 src/lib.rs                  | 15 +++++++++++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/examples/systray-example.rs b/examples/systray-example.rs
index 85c1a6e..b976b0a 100644
--- a/examples/systray-example.rs
+++ b/examples/systray-example.rs
@@ -33,7 +33,7 @@ fn main() -> Result<(), systray::Error> {
     })?;
 
     println!("Waiting on message!");
-    app.wait_for_message()?;
+    app.wait_for_message(None)?;
     Ok(())
 }
 
diff --git a/src/lib.rs b/src/lib.rs
index 36c6dbd..bb8f38a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -4,7 +4,8 @@ pub mod api;
 use std::{
     collections::HashMap,
     error, fmt,
-    sync::mpsc::{channel, Receiver},
+    sync::mpsc::{channel, Receiver, RecvTimeoutError},
+    time::Duration,
 };
 
 type BoxedError = Box<dyn error::Error + Send + Sync + 'static>;
@@ -133,12 +134,18 @@ impl Application {
         self.window.quit()
     }
 
-    pub fn wait_for_message(&mut self) -> Result<(), Error> {
+    ///
+    pub fn wait_for_message(&mut self, timeout: Option<Duration>) -> Result<(), Error> {
         loop {
             let msg;
-            match self.rx.recv() {
+            match self
+                .rx
+                .recv_timeout(timeout.unwrap_or_else(|| Duration::new(u64::MAX, 0)))
+            {
                 Ok(m) => msg = m,
-                Err(_) => {
+                // Yield and wait for the next poll
+                Err(RecvTimeoutError::Timeout) => break,
+                Err(RecvTimeoutError::Disconnected) => {
                     self.quit();
                     break;
                 }