[Proxying] Use Atomics.waitAsync to receive proxying notifications

emcc.py

@@ -1710,6 +1710,7 @@ def setup_pthreads(target):
     '__emscripten_thread_init',
     '__emscripten_thread_exit',
     '__emscripten_thread_crashed',
+    '__emscripten_thread_mailbox_await',
     '__emscripten_tls_init',
     '_pthread_self',
     'checkMailbox',

src/library_pthread.js

@@ -1219,18 +1219,40 @@ var LibraryPThread = {
   },
 #endif // MAIN_MODULE
 
-  $checkMailbox__deps: ['$callUserCallback'],
+  $checkMailbox__deps: ['$callUserCallback',
+                        '_emscripten_thread_mailbox_await'],
   $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()) {
+    var pthread_ptr = _pthread_self();
+    if (pthread_ptr) {
+      // If we are using Atomics.waitAsync as our notification mechanism, wait for
+      // a notification before processing the mailbox to avoid missing any work.
+      __emscripten_thread_mailbox_await(pthread_ptr);
       callUserCallback(() => __emscripten_check_mailbox());
     }
   },
 
-  _emscripten_notify_mailbox__deps: ['$checkMailbox'],
-  _emscripten_notify_mailbox__sig: 'vppp',
-  _emscripten_notify_mailbox: function(targetThreadId, currThreadId, mainThreadId) {
+  _emscripten_thread_mailbox_await__deps: ['$checkMailbox'],
+  _emscripten_thread_mailbox_await__sig: 'vp',
+  _emscripten_thread_mailbox_await: function(pthread_ptr) {
+    if (typeof Atomics.waitAsync === 'function') {
+      // TODO: How to make this work with wasm64?
+      var wait = Atomics.waitAsync(HEAP32, pthread_ptr >> 2, pthread_ptr);
+#if ASSERTIONS
+      assert(wait.async);
+#endif
+      wait.value.then(checkMailbox);
+      var waitingAsync = pthread_ptr + {{{ C_STRUCTS.pthread.waiting_async }}};
+      Atomics.store(HEAP32, waitingAsync >> 2, 1);
+    }
+  },
+
+  _emscripten_notify_mailbox_postmessage__deps: ['$checkMailbox'],
+  _emscripten_notify_mailbox_postmessage__sig: 'vppp',
+  _emscripten_notify_mailbox_postmessage: function(targetThreadId,
+                                                   currThreadId,
+                                                   mainThreadId) {
     if (targetThreadId == currThreadId) {
       setTimeout(() => checkMailbox());
     } else if (ENVIRONMENT_IS_PTHREAD) {
@@ -1241,7 +1263,7 @@ var LibraryPThread = {
 #if ASSERTIONS
         err('Cannot send message to thread with ID ' + targetThreadId + ', unknown thread ID!');
 #endif
-        return /*0*/;
+        return;
       }
       worker.postMessage({'cmd' : 'checkMailbox'});
     }

src/struct_info_internal.json

@@ -8,7 +8,8 @@
             "pthread": [
               "profilerBlock",
               "stack",
-              "stack_size"
+              "stack_size",
+              "waiting_async"
             ],
             "pthread_attr_t#": [
               "_a_transferredcanvases"

src/worker.js

@@ -216,6 +216,10 @@ function handleMessage(e) {
       // Pass the thread address to wasm to store it for fast access.
       Module['__emscripten_thread_init'](e.data.pthread_ptr, /*isMainBrowserThread=*/0, /*isMainRuntimeThread=*/0, /*canBlock=*/1);
 
+      // Await mailbox notifications with `Atomics.waitAsync` so we can start
+      // using the fast `Atomics.notify` notification path.
+      Module['__emscripten_thread_mailbox_await'](e.data.pthread_ptr);
+
 #if ASSERTIONS
       assert(e.data.pthread_ptr);
 #endif

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

@@ -97,6 +97,10 @@ struct pthread {
 	// wait until it reaches 0, at which point the mailbox is considered
 	// closed and no further messages will be enqueued.
 	_Atomic int mailbox_refcount;
+	// Whether the thread has executed a `waitAsync` on this pthread struct
+	// and can be notified of new mailbox messages via `Atomics.notify`.
+	// Otherwise the notification has to fall back to the postMessage path.
+	_Atomic int waiting_async;
 #endif
 #if _REENTRANT
 	_Atomic char sleeping;

system/lib/pthread/library_pthread.c

@@ -597,4 +597,5 @@ void __emscripten_init_main_thread(void) {
   __main_pthread.tsd = (void **)__pthread_tsd_main;
 
   _emscripten_thread_mailbox_init(&__main_pthread);
+  _emscripten_thread_mailbox_await(&__main_pthread);
 }

system/lib/pthread/thread_mailbox.c

@@ -70,6 +70,7 @@ void _emscripten_thread_mailbox_shutdown(pthread_t thread) {
 void _emscripten_thread_mailbox_init(pthread_t thread) {
   thread->mailbox = em_task_queue_create(thread);
   thread->mailbox_refcount = 1;
+  thread->waiting_async = 0;
 }
 
 // Exported for use in worker.js, but otherwise an internal function.
@@ -92,9 +93,9 @@ void _emscripten_check_mailbox() {
 // Send a postMessage notification telling the target thread to check its
 // mailbox when it returns to its event loop. Pass in the current thread and
 // main thread ids to minimize calls back into Wasm.
-void _emscripten_notify_mailbox(pthread_t target_thread,
-                                pthread_t curr_thread,
-                                pthread_t main_thread);
+void _emscripten_notify_mailbox_postmessage(pthread_t target_thread,
+                                            pthread_t curr_thread,
+                                            pthread_t main_thread);
 
 void emscripten_thread_mailbox_send(pthread_t thread, task t) {
   assert(thread->mailbox_refcount > 0);
@@ -112,8 +113,11 @@ void emscripten_thread_mailbox_send(pthread_t thread, task t) {
   notification_state previous =
     atomic_exchange(&thread->mailbox->notification, NOTIFICATION_PENDING);
   if (previous != NOTIFICATION_PENDING) {
-    _emscripten_notify_mailbox(thread,
-                               pthread_self(),
-                               emscripten_main_runtime_thread_id());
+    if (thread->waiting_async) {
+      __builtin_wasm_memory_atomic_notify((int*)thread, -1);
+    } else {
+      _emscripten_notify_mailbox_postmessage(
+        thread, pthread_self(), emscripten_main_runtime_thread_id());
+    }
   }
 }

system/lib/pthread/thread_mailbox.h

@@ -28,5 +28,7 @@ void emscripten_thread_mailbox_send(pthread_t thread, task t);
 // Initialize the mailbox on a pthread struct. Called during `pthread_create`.
 void _emscripten_thread_mailbox_init(pthread_t thread);
 
+void _emscripten_thread_mailbox_await(pthread_t thread);
+
 // Close the mailbox and cancel any pending messages.
 void _emscripten_thread_mailbox_shutdown(pthread_t thread);

system/lib/wasmfs/thread_utils.h

@@ -46,7 +46,8 @@ class ProxyWorker {
         // progress even if the main thread is blocked.
         //
         // TODO: Remove this once we can postMessage directly between workers
-        // without involving the main thread.
+        // without involving the main thread or once all browsers ship
+        // Atomics.waitAsync.
         //
         // Note that this requires adding _emscripten_proxy_execute_queue to
         // EXPORTED_FUNCTIONS.

test/other/metadce/test_metadce_minimal_pthreads.exports

@@ -3,7 +3,7 @@ B
 C
 D
 E
-o
+F
 p
 q
 r

test/other/metadce/test_metadce_minimal_pthreads.imports

@@ -12,3 +12,4 @@ a.k
 a.l
 a.m
 a.n
+a.o

test/other/metadce/test_metadce_minimal_pthreads.jssize

@@ -1 +1 @@
-15320
+15563

test/other/metadce/test_metadce_minimal_pthreads.sent

@@ -12,3 +12,4 @@ k
 l
 m
 n
+o

test/other/metadce/test_metadce_minimal_pthreads.size

@@ -1 +1 @@
-18922
+18954

test/pthread/test_pthread_proxying_refcount.c

@@ -50,18 +50,6 @@ void* execute_and_free_queue(void* arg) {
     em_proxying_queue_destroy(queues[i]);
   }
 
-  // Wrap the normal worker event listener so that we can determine when our
-  // proxying events have been received and handled.
-  EM_ASM({
-      var oldOnMessage = onmessage;
-      onmessage = (e) => {
-        oldOnMessage(e);
-        if (e.data.cmd == 'checkMailbox') {
-          _register_processed();
-        }
-      };
-    });
-
   // Exit with a live runtime so the queued work notification is received and we
   // try to execute the queue again, even though we already executed all its
   // work and we are now just waiting for the notifications to be received so we
@@ -90,9 +78,12 @@ int main() {
   while (!executed[0] || !executed[1]) {
   }
 
-  // Wait for the postMessage notification to be received.
-  while (processed < 1) {
-  }
+  // Wait a bit (20 ms) for the notification to be received.
+  struct timespec time = {
+    .tv_sec = 0,
+    .tv_nsec = 20 * 1000 * 1000,
+  };
+  nanosleep(&time, NULL);
 
 #ifndef SANITIZER
   // Our zombies should not have been freed yet.

test/reference_struct_info.json

@@ -1351,10 +1351,11 @@
             "p_proto": 8
         },
         "pthread": {
-            "__size__": 128,
+            "__size__": 132,
             "profilerBlock": 112,
             "stack": 52,
-            "stack_size": 56
+            "stack_size": 56,
+            "waiting_async": 128
         },
         "pthread_attr_t": {
             "__size__": 44,