[Proxying] Send messages via in-memory mailbox queues (#18852)

Threads were previously notified of new work via postMessage messages that
carried pointers to the task queues to execute. There was no way to synchronously
pump or inspect these pending messages however, and there is no central registry
of all task queues for a thread, so this mechanism afforded no way to discover
or cancel pending work when a thread dies.

In preparation for implementing work cancellation, move the pending messages
into userspace by giving each thread a "mailbox", which is an `em_task_queue` in
the pthread struct. Instead of using `postMessage`, proxying queues now use the
thread mailbox API to notify threads of new work.

Internally, thread mailboxes still use postMessage to schedule work to be
executed when a thread returns to its event loop. Since the only task queues
involved in postMessages are now at known locations relative to the pthread
struct, there is no longer any need to store pointers to them in the postMessage
messages themselves. Removing these pointers works around tricky notification
and lifetime management edge cases that would have caused problems such as
dropped work or use-after-free bugs in future PRs.

When a thread dies because it exits or is canceled, it "closes" its mailbox by
decrementing a refcount and waiting to observe a refcount of 0. At this point,
the thread mailbox API ensures that no new messages will be enqueued on the
mailbox. Because the postMessage messages no longer contain task queue pointers,
it is safe to destroy the mailbox immediately after it is closed.

Author: tlively

emcc.py

@@ -1712,7 +1712,7 @@ def setup_pthreads(target):
     '__emscripten_thread_crashed',
     '__emscripten_tls_init',
     '_pthread_self',
-    'executeNotifiedProxyingQueue',
+    'checkMailbox',
   ]
   settings.EXPORTED_FUNCTIONS += worker_imports
   building.user_requested_exports.update(worker_imports)

src/library_pthread.js

@@ -270,8 +270,8 @@ var LibraryPThread = {
           return;
         }
 
-        if (cmd === 'processProxyingQueue') {
-          executeNotifiedProxyingQueue(d['queue']);
+        if (cmd === 'checkMailbox') {
+          checkMailbox();
         } else if (cmd === 'spawnThread') {
           spawnThread(d);
         } else if (cmd === 'cleanupThread') {
@@ -1212,29 +1212,22 @@ var LibraryPThread = {
   },
 #endif // MAIN_MODULE
 
-  $executeNotifiedProxyingQueue__deps: ['$callUserCallback'],
-  $executeNotifiedProxyingQueue: function(queue) {
-    // Set the notification state to processing.
-    Atomics.store(HEAP32, queue >> 2, {{{ cDefine('NOTIFICATION_RECEIVED') }}});
-    // Only execute the queue if we have a live pthread runtime. We
-    // implement pthread_self to return 0 if there is no live runtime.
+  $checkMailbox__deps: ['$callUserCallback'],
+  $checkMailbox: function() {
+    // Only check the mailbox if we have a live pthread runtime. We implement
+    // pthread_self to return 0 if there is no live runtime.
     if (_pthread_self()) {
-      callUserCallback(() => __emscripten_proxy_execute_task_queue(queue));
+      callUserCallback(() => __emscripten_check_mailbox());
     }
-    // Set the notification state to none as long as a new notification has not
-    // been sent while we were processing.
-    Atomics.compareExchange(HEAP32, queue >> 2,
-                            {{{ cDefine('NOTIFICATION_RECEIVED') }}},
-                            {{{ cDefine('NOTIFICATION_NONE') }}});
   },
 
-  _emscripten_notify_task_queue__deps: ['$executeNotifiedProxyingQueue'],
-  _emscripten_notify_task_queue__sig: 'vpppp',
-  _emscripten_notify_task_queue: function(targetThreadId, currThreadId, mainThreadId, queue) {
+  _emscripten_notify_mailbox__deps: ['$checkMailbox'],
+  _emscripten_notify_mailbox__sig: 'vppp',
+  _emscripten_notify_mailbox: function(targetThreadId, currThreadId, mainThreadId) {
     if (targetThreadId == currThreadId) {
-      setTimeout(() => executeNotifiedProxyingQueue(queue));
+      setTimeout(() => checkMailbox());
     } else if (ENVIRONMENT_IS_PTHREAD) {
-      postMessage({'targetThread' : targetThreadId, 'cmd' : 'processProxyingQueue', 'queue' : queue});
+      postMessage({'targetThread' : targetThreadId, 'cmd' : 'checkMailbox'});
     } else {
       var worker = PThread.pthreads[targetThreadId];
       if (!worker) {
@@ -1243,7 +1236,7 @@ var LibraryPThread = {
 #endif
         return /*0*/;
       }
-      worker.postMessage({'cmd' : 'processProxyingQueue', 'queue': queue});
+      worker.postMessage({'cmd' : 'checkMailbox'});
     }
   }
 };

src/worker.js

@@ -52,10 +52,6 @@ if (ENVIRONMENT_IS_NODE) {
 // Thread-local guard variable for one-time init of the JS state
 var initializedJS = false;
 
-// Proxying queues that were notified before the thread started and need to be
-// executed as part of startup.
-var pendingNotifiedProxyingQueues = [];
-
 #if ASSERTIONS
 function assert(condition, text) {
   if (!condition) abort('Assertion failed: ' + text);
@@ -237,15 +233,6 @@ function handleMessage(e) {
         // We only do this once per worker since they get reused
         Module['__embind_initialize_bindings']();
 #endif // EMBIND
-
-        // Execute any proxied work that came in before the thread was
-        // initialized. Only do this once because it is only possible for
-        // proxying notifications to arrive before thread initialization on
-        // fresh workers.
-        pendingNotifiedProxyingQueues.forEach(queue => {
-          Module['executeNotifiedProxyingQueue'](queue);
-        });
-        pendingNotifiedProxyingQueues = [];
         initializedJS = true;
       }
 
@@ -268,12 +255,9 @@ function handleMessage(e) {
       }
     } else if (e.data.target === 'setimmediate') {
       // no-op
-    } else if (e.data.cmd === 'processProxyingQueue') {
+    } else if (e.data.cmd === 'checkMailbox') {
       if (initializedJS) {
-        Module['executeNotifiedProxyingQueue'](e.data.queue);
-      } else {
-        // Defer executing this queue until the runtime is initialized.
-        pendingNotifiedProxyingQueues.push(e.data.queue);
+        Module['checkMailbox']();
       }
     } else if (e.data.cmd) {
       // The received message looks like something that should be handled by this message

system/lib/libc/musl/src/internal/pthread_impl.h

@@ -10,8 +10,10 @@
 #include "syscall.h"
 #include "atomic.h"
 #ifdef __EMSCRIPTEN__
-#include <emscripten/threading.h>
+#include "em_task_queue.h"
+#include "thread_mailbox.h"
 #include "threading_internal.h"
+#include <emscripten/threading.h>
 #endif
 #include "futex.h"
 
@@ -78,6 +80,23 @@ struct pthread {
 	// The TLS base to use the main module TLS data.  Secondary modules
 	// still require dynamic allocation.
 	void* tls_base;
+	// The lowest level of the proxying system. Other threads can enqueue
+	// messages on the mailbox and notify this thread to asynchronously
+	// process them once it returns to its event loop. When this thread is
+	// shut down, the mailbox is closed (see below) to prevent further
+	// messages from being enqueued and all the remaining queued messages
+	// are dequeued and their shutdown handlers are executed. This allows
+	// other threads waiting for their messages to be processed to be
+	// notified that their messages will not be processed after all.
+	em_task_queue* mailbox;
+	// To ensure that no other thread is concurrently enqueueing a message
+	// when this thread shuts down, maintain an atomic refcount. Enqueueing
+	// threads atomically increment the count from a nonzero number to
+	// acquire the mailbox and decrement the count when they finish. When
+	// this thread shuts down it will atomically decrement the count and
+	// wait until it reaches 0, at which point the mailbox is considered
+	// closed and no further messages will be enqueued.
+	_Atomic int mailbox_refcount;
 #endif
 #if _REENTRANT
 	_Atomic char sleeping;

system/lib/pthread/em_task_queue.c

@@ -13,6 +13,7 @@
 
 #include "em_task_queue.h"
 #include "proxying_notification_state.h"
+#include "thread_mailbox.h"
 
 #define EM_TASK_QUEUE_INITIAL_CAPACITY 128
 
@@ -194,26 +195,41 @@ task em_task_queue_dequeue(em_task_queue* queue) {
   return t;
 }
 
-// Send a postMessage notification containing the em_task_queue pointer to the
-// target thread so it will execute the queue when it returns to the event loop.
-// Also pass in the current thread and main thread ids to minimize calls back
-// into Wasm.
-void _emscripten_notify_task_queue(pthread_t target_thread,
-                                   pthread_t curr_thread,
-                                   pthread_t main_thread,
-                                   em_task_queue* queue);
-
-void em_task_queue_notify(em_task_queue* queue) {
-  // If there is no pending notification for this queue, create one. If an old
-  // notification is currently being processed, it may or may not execute this
-  // work. In case it does not, the new notification will ensure the work is
-  // still executed.
+static void receive_notification(void* arg) {
+  em_task_queue* tasks = arg;
+  tasks->notification = NOTIFICATION_RECEIVED;
+  em_task_queue_execute(tasks);
+  notification_state expected = NOTIFICATION_RECEIVED;
+  atomic_compare_exchange_strong(
+    &tasks->notification, &expected, NOTIFICATION_NONE);
+}
+
+int em_task_queue_send(em_task_queue* queue, task t) {
+  // Ensure the target mailbox will remain open or detect that it is already
+  // closed.
+  if (!emscripten_thread_mailbox_ref(queue->thread)) {
+    return 0;
+  }
+
+  pthread_mutex_lock(&queue->mutex);
+  int enqueued = em_task_queue_enqueue(queue, t);
+  pthread_mutex_unlock(&queue->mutex);
+  if (!enqueued) {
+    emscripten_thread_mailbox_unref(queue->thread);
+    return 0;
+  }
+
+  // We're done if there is already a pending notification for this task queue.
+  // Otherwise, we will send one.
   notification_state previous =
     atomic_exchange(&queue->notification, NOTIFICATION_PENDING);
-  if (previous != NOTIFICATION_PENDING) {
-    _emscripten_notify_task_queue(queue->thread,
-                                  pthread_self(),
-                                  emscripten_main_runtime_thread_id(),
-                                  queue);
+  if (previous == NOTIFICATION_PENDING) {
+    emscripten_thread_mailbox_unref(queue->thread);
+    return 1;
   }
+
+  emscripten_thread_mailbox_send(
+    queue->thread, (task){.func = receive_notification, .arg = queue});
+  emscripten_thread_mailbox_unref(queue->thread);
+  return 1;
 }

system/lib/pthread/em_task_queue.h

@@ -50,7 +50,7 @@ em_task_queue* em_task_queue_create(pthread_t thread);
 
 void em_task_queue_destroy(em_task_queue* queue);
 
-// Execute tasks until an empty queue is observed.
+// Execute tasks until an empty queue is observed. Internally locks the queue.
 void em_task_queue_execute(em_task_queue* queue);
 
 // Not thread safe.
@@ -69,6 +69,7 @@ int em_task_queue_enqueue(em_task_queue* queue, task t);
 // Not thread safe. Assumes the queue is not empty.
 task em_task_queue_dequeue(em_task_queue* queue);
 
-// Schedule the queue to be executed next time its owning thread returns to its
-// event loop.
-void em_task_queue_notify(em_task_queue* queue);
+// Atomically enqueue the task and schedule the queue to be executed next time
+// its owning thread returns to its event loop. Returns 1 on success and 0
+// otherwise. Internally locks the queue.
+int em_task_queue_send(em_task_queue* queue, task t);

system/lib/pthread/library_pthread.c

@@ -595,4 +595,6 @@ void __emscripten_init_main_thread(void) {
   // this is used by pthread_key_delete for deleting thread-specific data.
   __main_pthread.next = __main_pthread.prev = &__main_pthread;
   __main_pthread.tsd = (void **)__pthread_tsd_main;
+
+  _emscripten_thread_mailbox_init(&__main_pthread);
 }

system/lib/pthread/proxying.c

@@ -9,11 +9,12 @@
 #include <emscripten/proxying.h>
 #include <emscripten/threading.h>
 #include <pthread.h>
+#include <stdatomic.h>
 #include <stdlib.h>
 #include <string.h>
 
 #include "em_task_queue.h"
-#include "proxying_notification_state.h"
+#include "thread_mailbox.h"
 
 struct em_proxying_queue {
   // Protects all accesses to em_task_queues, size, and capacity.
@@ -103,17 +104,6 @@ static em_task_queue* get_or_add_tasks_for_thread(em_proxying_queue* q,
   return tasks;
 }
 
-// Exported for use in worker.js, but otherwise an internal function.
-EMSCRIPTEN_KEEPALIVE
-void _emscripten_proxy_execute_task_queue(em_task_queue* tasks) {
-  // Before we attempt to execute a request from another thread make sure we
-  // are in sync with all the loaded code.
-  // For example, in PROXY_TO_PTHREAD the atexit functions are called via
-  // a proxied call, and without this call to syncronize we would crash if
-  // any atexit functions were registered from a side module.
-  em_task_queue_execute(tasks);
-}
-
 void emscripten_proxy_execute_queue(em_proxying_queue* q) {
   assert(q != NULL);
   assert(pthread_self());
@@ -156,15 +146,8 @@ int emscripten_proxy_async(em_proxying_queue* q,
   if (tasks == NULL) {
     return 0;
   }
-  pthread_mutex_lock(&tasks->mutex);
-  int enqueued = em_task_queue_enqueue(tasks, (task){func, arg});
-  pthread_mutex_unlock(&tasks->mutex);
-  if (!enqueued) {
-    return 0;
-  }
 
-  em_task_queue_notify(tasks);
-  return 1;
+  return em_task_queue_send(tasks, (task){func, arg});
 }
 
 struct em_proxying_ctx {

system/lib/pthread/pthread_create.c

@@ -223,6 +223,8 @@ int __pthread_create(pthread_t* restrict res,
   _emscripten_thread_profiler_init(new);
 #endif
 
+  _emscripten_thread_mailbox_init(new);
+
   struct pthread *self = __pthread_self();
   dbg("start __pthread_create: new=%p new_end=%p stack=%p->%p "
       "stack_size=%zu tls_base=%p",
@@ -303,6 +305,8 @@ void _emscripten_thread_exit(void* result) {
   self->cancelasync = PTHREAD_CANCEL_DEFERRED;
   self->result = result;
 
+  _emscripten_thread_mailbox_shutdown(self);
+
   // Run any handlers registered with pthread_cleanup_push
   __run_cleanup_handlers();