1 files changed, 20 insertions, 5 deletions

diff --git a/ot_mutex.c b/ot_mutex.c
index 37d9962..5392264 100644
--- a/ot_mutex.c
+++ b/ot_mutex.c
@@ -6,6 +6,7 @@
 /* System */
 #include <arpa/inet.h>
 #include <pthread.h>
+#include <stdatomic.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <unistd.h>
@@ -25,7 +26,9 @@
 /* Our global all torrents list */
 static ot_vector       all_torrents[OT_BUCKET_COUNT];
 static pthread_mutex_t bucket_mutex[OT_BUCKET_COUNT];
-static size_t          g_torrent_count;
+/* Every bucket's unlock adjusts this, and a bucket lock only keeps other
+   threads out of that one bucket, so the counter needs its own protection */
+static _Atomic size_t  g_torrent_count;
 
 /* Self pipe from opentracker.c */
 extern int             g_self_pipe[2];
@@ -47,14 +50,15 @@ ot_vector *mutex_bucket_lock_by_hash(ot_hash const hash) { return mutex_bucket_l
 
 void mutex_bucket_unlock(int bucket, int delta_torrentcount) {
   pthread_mutex_unlock(bucket_mutex + bucket);
-  g_torrent_count += delta_torrentcount;
+  if (delta_torrentcount)
+    atomic_fetch_add_explicit(&g_torrent_count, (size_t)delta_torrentcount, memory_order_relaxed);
 }
 
 void mutex_bucket_unlock_by_hash(ot_hash const hash, int delta_torrentcount) {
   mutex_bucket_unlock(bucket_from_hash(hash), delta_torrentcount);
 }
 
-size_t mutex_get_torrent_count(void) { return g_torrent_count; }
+size_t mutex_get_torrent_count(void) { return atomic_load_explicit(&g_torrent_count, memory_order_relaxed); }
 
 /* TaskQueue Magic */
 
@@ -145,9 +149,13 @@ ot_taskid mutex_workqueue_poptask(ot_tasktype *tasktype) {
         break;
       }
 
-    /* Wait until the next task is being fed */
-    if (!taskid)
+    /* Wait until the next task is being fed, unless the tracker is on its
+       way out, in which case the worker is meant to leave rather than wait */
+    if (!taskid) {
+      if (!g_opentracker_running)
+        break;
       pthread_cond_wait(&tasklist_being_filled, &tasklist_mutex);
+    }
   }
 
   /* Release lock */
@@ -271,6 +279,13 @@ void mutex_init(void) {
   memset(all_torrents, 0, sizeof(all_torrents));
 }
 
+/* Let every worker parked on the queue see that the tracker is stopping */
+void mutex_workqueue_wakeall(void) {
+  pthread_mutex_lock(&tasklist_mutex);
+  pthread_cond_broadcast(&tasklist_being_filled);
+  pthread_mutex_unlock(&tasklist_mutex);
+}
+
 void mutex_deinit(void) {
   size_t i;
   for (i = 0; i < OT_BUCKET_COUNT; ++i)