From be5ee9f7934a9a8b8443c9590bc45c41ebb4a844 Mon Sep 17 00:00:00 2001
From: pixl <pixlone@protonmail.com>
Date: Fri, 12 May 2023 17:25:28 -0400
Subject: [PATCH] Exit gracefully upon return

Stop all worker threads when main thread exits.
---
 src/logid/util/task.cpp | 33 +++++++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/logid/util/task.cpp b/src/logid/util/task.cpp
index fea2f5e..009594d 100644
--- a/src/logid/util/task.cpp
+++ b/src/logid/util/task.cpp
@@ -18,6 +18,7 @@
 #include <util/task.h>
 #include <queue>
 #include <optional>
+#include <cassert>
 
 using namespace logid;
 using namespace std::chrono;
@@ -35,18 +36,38 @@ static std::priority_queue<task, std::vector<task>, task_less> tasks {};
 static std::mutex task_mutex {};
 static std::condition_variable task_cv {};
 static std::atomic_bool workers_init = false;
+static std::atomic_bool workers_run = false;
 
-[[noreturn]] static void worker() {
+void stop_workers() {
     std::unique_lock lock(task_mutex);
-    while (true) {
-        task_cv.wait(lock, []() { return !tasks.empty(); });
+    if (workers_init) {
+        workers_run = false;
+        lock.unlock();
+        task_cv.notify_all();
+
+        /* Wait for all workers to end */
+        lock.lock();
+    }
+}
+
+void worker() {
+    std::unique_lock lock(task_mutex);
+    while (workers_run) {
+        task_cv.wait(lock, []() { return !tasks.empty() || !workers_run; });
+
+        if (!workers_run)
+            break;
 
         /* top task is in the future, wait */
         if (tasks.top().time >= system_clock::now()) {
             auto wait = tasks.top().time - system_clock::now();
             task_cv.wait_for(lock, wait, []() {
-                return !tasks.empty() && (tasks.top().time < system_clock::now());
+                return (!tasks.empty() && (tasks.top().time < system_clock::now())) ||
+                    !workers_run;
             });
+
+            if (!workers_run)
+                break;
         }
 
         if (!tasks.empty()) {
@@ -67,11 +88,15 @@ static std::atomic_bool workers_init = false;
 
 void logid::init_workers(int worker_count) {
     std::lock_guard lock(task_mutex);
+    assert(!workers_init);
 
     for (int i = 0; i < worker_count; ++i)
         std::thread(&worker).detach();
 
     workers_init = true;
+    workers_run = true;
+
+    atexit(&stop_workers);
 }
 
 void logid::run_task(std::function<void()> function) {