From 12121bae375d33585e18d0a7f2efcd9a9998e0d1 Mon Sep 17 00:00:00 2001
From: Firstyear <william@blackhats.net.au>
Date: Tue, 4 Jul 2023 15:53:48 +1000
Subject: [PATCH] Improve tasks daemon shutdown (#1806)

---
 unix_integration/src/tasks_daemon.rs | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/unix_integration/src/tasks_daemon.rs b/unix_integration/src/tasks_daemon.rs
index 5bc2338d7..b596cb5c4 100644
--- a/unix_integration/src/tasks_daemon.rs
+++ b/unix_integration/src/tasks_daemon.rs
@@ -330,6 +330,7 @@ async fn main() -> ExitCode {
             debug!("Attempting to use {} ...", task_sock_path);
 
             let (broadcast_tx, mut broadcast_rx) = broadcast::channel(4);
+            let mut d_broadcast_rx = broadcast_tx.subscribe();
 
             let server = tokio::spawn(async move {
                 loop {
@@ -345,7 +346,14 @@ async fn main() -> ExitCode {
                                     info!("Found kanidm_unixd, waiting for tasks ...");
                                     // Yep! Now let the main handler do it's job.
                                     // If it returns (dc, etc, then we loop and try again).
-                                    handle_tasks(stream, &cfg).await;
+                                    tokio::select! {
+                                        _ = d_broadcast_rx.recv() => {
+                                            break;
+                                        }
+                                        _ = handle_tasks(stream, &cfg) => {
+                                            continue;
+                                        }
+                                    }
                                 }
                                 Err(e) => {
                                     debug!("\\---> {:?}", e);