Do not close connections immediately on close signal, await for remaining responses

1 files changed, 20 insertions, 0 deletions

diff --git a/src/util.rs b/src/util.rs
index f09a3bc..017ef00 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -1,5 +1,7 @@
 use serde::Serialize;
 
+use tokio::sync::watch;
+
 /// Utility function: encodes any serializable value in MessagePack binary format
 /// using the RMP library.
 ///
@@ -16,3 +18,21 @@ where
 	val.serialize(&mut se)?;
 	Ok(wr)
 }
+
+/// This async function returns only when a true signal was received
+/// from a watcher that tells us when to exit.
+/// Usefull in a select statement to interrupt another
+/// future:
+/// ```
+/// select!(
+///     _ = a_long_task() => Success,
+///     _ = await_exit(must_exit) => Interrupted,
+/// )
+/// ```
+pub async fn await_exit(mut must_exit: watch::Receiver<bool>) {
+	loop {
+		if must_exit.recv().await == Some(true) {
+			return;
+		}
+	}
+}