Fix a concurrency bug in OP_WAIT_TIMEOUT

src/libAtomVM/jit.c

@@ -888,6 +888,7 @@ static Context *jit_process_signal_messages(Context *ctx, JITState *jit_state)
 static term jit_mailbox_peek(Context *ctx)
 {
     TRACE("jit_mailbox_peek: ctx->process_id=%" PRId32 "\n", ctx->process_id);
+    ctx->mailbox.receive_has_match_clauses = true;
     term out = term_invalid_term();
     mailbox_peek(ctx, &out);
     return out;
@@ -897,12 +898,14 @@ static void jit_mailbox_remove_message(Context *ctx)
 {
     TRACE("jit_mailbox_remove_message: ctx->process_id=%" PRId32 "\n", ctx->process_id);
     mailbox_remove_message(&ctx->mailbox, &ctx->heap);
+    ctx->mailbox.receive_has_match_clauses = false;
 }
 
 static void jit_timeout(Context *ctx)
 {
     TRACE("jit_timeout: ctx->process_id=%" PRId32 "\n", ctx->process_id);
     context_update_flags(ctx, ~WaitingTimeoutExpired, NoFlags);
+    ctx->mailbox.receive_has_match_clauses = false;
     mailbox_reset(&ctx->mailbox);
 }
 
@@ -952,18 +955,23 @@ static Context *jit_wait_timeout(Context *ctx, JITState *jit_state, term timeout
         needs_to_wait = 1;
     } else if (context_get_flags(ctx, WaitingTimeout) != 0) {
         needs_to_wait = 1;
-    } else if (!mailbox_has_next(&ctx->mailbox)) {
-        needs_to_wait = 1;
     }
+    // else: WaitingTimeoutExpired -- fall through to timeout.
+    // Any messages in the mailbox are left for the next receive.
 
     if (needs_to_wait) {
+        // Signal processing may have moved messages to the inner list.
+        // If there are match clauses (loop_rec was executed), jump to
+        // loop_rec to scan them.
+        if (ctx->mailbox.receive_has_match_clauses && mailbox_has_next(&ctx->mailbox)) {
+            jit_state->continuation = module_get_native_entry_point(jit_state->module, label);
+            return ctx;
+        }
         return jit_schedule_wait_cp(ctx, jit_state);
-    } else {
-        // clang cannot tail-optimize this, so return to loop to avoid any stack overflow
-        // __attribute__((musttail)) return jit_state->continuation(ctx, jit_state, &module_native_interface);
-        jit_state->continuation = module_get_native_entry_point(jit_state->module, label);
-        return ctx;
     }
+    // else: timer expired, fall through to timeout
+    // jit_state->continuation already points to the trap handler code
+    return ctx;
 }
 
 static Context *jit_wait_timeout_trap_handler(Context *ctx, JITState *jit_state, int label)
@@ -976,6 +984,7 @@ static Context *jit_wait_timeout_trap_handler(Context *ctx, JITState *jit_state,
         return scheduler_wait(ctx);
     }
 
+    // Messages available, jump to loop_rec to scan them.
     jit_state->continuation = module_get_native_entry_point(jit_state->module, label);
     return ctx;
 }

src/libAtomVM/mailbox.c

@@ -45,6 +45,7 @@ void mailbox_init(Mailbox *mbx)
     mbx->inner_last = NULL;
     mbx->receive_pointer = NULL;
     mbx->receive_pointer_prev = NULL;
+    mbx->receive_has_match_clauses = false;
 }
 
 // Convert a mailbox message (struct Message or struct TermSignal) to a heap

src/libAtomVM/mailbox.h

@@ -178,6 +178,10 @@ typedef struct
     // Receive pointers are on inner list items.
     MailboxMessage *receive_pointer;
     MailboxMessage *receive_pointer_prev;
+    // Set by loop_rec, cleared by remove_message and timeout.
+    // Used by wait_timeout to know if there are match clauses to try
+    // when signal processing moved messages to the inner list.
+    bool receive_has_match_clauses;
 } Mailbox;
 
 // TODO: a lot of this code depends on Context * and should be decoupled

src/libAtomVM/opcodesswitch.h

@@ -2798,6 +2798,7 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                     }
                     PROCESS_SIGNAL_MESSAGES();
                     mailbox_remove_message(&ctx->mailbox, &ctx->heap);
+                    ctx->mailbox.receive_has_match_clauses = false;
                     // Cannot GC now as remove_message is GC neutral
                 #endif
                 break;
@@ -2808,6 +2809,7 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
 
                 #ifdef IMPL_EXECUTE_LOOP
                     context_update_flags(ctx, ~WaitingTimeoutExpired, NoFlags);
+                    ctx->mailbox.receive_has_match_clauses = false;
 
                     mailbox_reset(&ctx->mailbox);
                 #endif
@@ -2823,6 +2825,7 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                 TRACE("loop_rec/2, dreg=%c%i\n", T_DEST_REG(dreg));
 
                 #ifdef IMPL_EXECUTE_LOOP
+                    ctx->mailbox.receive_has_match_clauses = true;
                     term ret;
                     PROCESS_SIGNAL_MESSAGES();
                     if (mailbox_peek(ctx, &ret)) {
@@ -2898,15 +2901,19 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                         needs_to_wait = 1;
                     } else if (context_get_flags(ctx, WaitingTimeout) != 0) {
                         needs_to_wait = 1;
-                    } else if (!mailbox_has_next(&ctx->mailbox)) {
-                        needs_to_wait = 1;
                     }
+                    // else: WaitingTimeoutExpired -- fall through to timeout.
+                    // Any messages in the mailbox are left for the next receive.
 
                     if (needs_to_wait) {
+                        // Signal processing may have moved messages to the inner
+                        // list. If there are match clauses (loop_rec was
+                        // executed), jump to loop_rec to scan them.
+                        if (ctx->mailbox.receive_has_match_clauses && mailbox_has_next(&ctx->mailbox)) {
+                            JUMP_TO_ADDRESS(mod->labels[label]);
+                        }
                         ctx->waiting_with_timeout = true;
                         SCHEDULE_WAIT(mod, saved_pc);
-                    } else {
-                        JUMP_TO_ADDRESS(mod->labels[label]);
                     }
                 #endif
 
@@ -2922,8 +2929,6 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
 #ifdef IMPL_EXECUTE_LOOP
 wait_timeout_trap_handler:
             {
-                // Determine if a message arrived to either jump to timeout label
-                // or to continuation.
                 // Redo the offset computation and refetch the label
                 int label;
                 DECODE_LABEL(label, pc)
@@ -2932,6 +2937,8 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                 TRACE("wait_timeout_trap_handler, label: %i\n", label);
                 PROCESS_SIGNAL_MESSAGES();
                 if (context_get_flags(ctx, WaitingTimeoutExpired)) {
+                    // Timer expired -- fall through to timeout.
+                    // Any messages in the mailbox are left for the next receive.
                     ctx->waiting_with_timeout = false;
                 } else {
                     if (UNLIKELY(!mailbox_has_next(&ctx->mailbox))) {