rts: Add missing write barriers in MVar wake-up paths

Previously PerformPut failed to respect the non-moving collector's
snapshot invariant, hiding references to an MVar and its new value by
overwriting a stack frame without dirtying the stack. Fix this.

PerformTake exhibited a similar bug, failing to dirty (and therefore
mark) the blocked stack before mutating it.

Closes #20399.

2 files changed, 37 insertions, 12 deletions

diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index 8f99105b18..26c4cfe7b4 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -1591,6 +1591,23 @@ stg_writeTVarzh (P_ tvar,     /* :: TVar a */
  * exception and never perform its take or put, and we'd end up with a
  * deadlock.
  *
+ * Note [Nonmoving write barrier in Perform{Take,Put}]
+ * ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ * As noted in Note [Non-moving garbage collector] in NonMoving.c, the
+ * non-moving GC requires that all overwritten pointers be pushed to the update
+ * remembered set. In the case of stack mutation this typically happens by
+ * "dirtying" the stack, which eagerly traces the entire stack chunk.
+ *
+ * An exception to this rule is PerformPut, which mutates the stack of a
+ * blocked thread (overwriting an stg_block_putmvar frame). To ensure that the
+ * collector sees the MVar and value reachable from the overwritten frame, we
+ * must push them to the update remembered set. Failing to do so was the cause
+ * of #20399.
+ *
+ * Note that unlike PerformPut, the callers of PerformTake first dirty the
+ * stack prior mutating it (since they introduce a *new*, potentially
+ * inter-generational reference to the stack) and therefore the barrier
+ * described above is unnecessary in this case.
  * -------------------------------------------------------------------------- */
 
 stg_isEmptyMVarzh ( P_ mvar /* :: MVar a */ )
@@ -1619,15 +1636,22 @@ stg_newMVarzh ()
 }
 
 
+// See Note [Nonmoving write barrier in Perform{Put,Take}].
+// Precondition: the stack must be dirtied.
 #define PerformTake(stack, value)               \
     W_ sp;                                      \
     sp = StgStack_sp(stack);                    \
     W_[sp + WDS(1)] = value;                    \
     W_[sp + WDS(0)] = stg_ret_p_info;
 
+// See Note [Nonmoving write barrier in Perform{Put,Take}].
 #define PerformPut(stack,lval)                  \
     W_ sp;                                      \
     sp = StgStack_sp(stack) + WDS(3);           \
+    IF_NONMOVING_WRITE_BARRIER_ENABLED {        \
+      ccall updateRemembSetPushClosure_(BaseReg "ptr", W_[sp - WDS(1)] "ptr"); \
+      ccall updateRemembSetPushClosure_(BaseReg "ptr", W_[sp - WDS(2)] "ptr"); \
+    }                                           \
     StgStack_sp(stack) = sp;                    \
     lval = W_[sp - WDS(1)];
 
@@ -1900,15 +1924,14 @@ loop:
     // actually perform the takeMVar
     W_ stack;
     stack = StgTSO_stackobj(tso);
+    if ((TO_W_(StgStack_dirty(stack)) & STACK_DIRTY) == 0) {
+        ccall dirty_STACK(MyCapability() "ptr", stack "ptr");
+    }
     PerformTake(stack, val);
 
     // indicate that the MVar operation has now completed.
     StgTSO__link(tso) = stg_END_TSO_QUEUE_closure;
 
-    if ((TO_W_(StgStack_dirty(stack)) & STACK_DIRTY) == 0) {
-        ccall dirty_STACK(MyCapability() "ptr", stack "ptr");
-    }
-
     ccall tryWakeupThread(MyCapability() "ptr", tso);
 
     // If it was a readMVar, then we can still do work,
@@ -1991,15 +2014,14 @@ loop:
     // actually perform the takeMVar
     W_ stack;
     stack = StgTSO_stackobj(tso);
+    if ((TO_W_(StgStack_dirty(stack)) & STACK_DIRTY) == 0) {
+        ccall dirty_STACK(MyCapability() "ptr", stack "ptr");
+    }
     PerformTake(stack, val);
 
     // indicate that the MVar operation has now completed.
     StgTSO__link(tso) = stg_END_TSO_QUEUE_closure;
 
-    if ((TO_W_(StgStack_dirty(stack)) & STACK_DIRTY) == 0) {
-        ccall dirty_STACK(MyCapability() "ptr", stack "ptr");
-    }
-
     ccall tryWakeupThread(MyCapability() "ptr", tso);
 
     // If it was a readMVar, then we can still do work,
@@ -2273,15 +2295,14 @@ loop:
     // actually perform the takeMVar
     W_ stack;
     stack = StgTSO_stackobj(tso);
+    if (TO_W_(StgStack_dirty(stack)) == 0) {
+        ccall dirty_STACK(MyCapability() "ptr", stack "ptr");
+    }
     PerformTake(stack, val);
 
     // indicate that the operation has now completed.
     StgTSO__link(tso) = stg_END_TSO_QUEUE_closure;
 
-    if (TO_W_(StgStack_dirty(stack)) == 0) {
-        ccall dirty_STACK(MyCapability() "ptr", stack "ptr");
-    }
-
     ccall tryWakeupThread(MyCapability() "ptr", tso);
 
     // For MVars we loop here, waking up all readers.
diff --git a/rts/sm/NonMoving.c b/rts/sm/NonMoving.c
index b92bf2276e..dd019ec18b 100644
--- a/rts/sm/NonMoving.c
+++ b/rts/sm/NonMoving.c
@@ -229,6 +229,10 @@ Mutex concurrent_coll_finished_lock;
  *  - Note [StgStack dirtiness flags and concurrent marking] (TSO.h) describes
  *    the protocol for concurrent marking of stacks.
  *
+ *  - Note [Nonmoving write barrier in Perform{Take,Put}] (PrimOps.cmm) describes
+ *    a tricky barrier necessary when resuming threads blocked on MVar
+ *    operations.
+ *
  *  - Note [Static objects under the nonmoving collector] (Storage.c) describes
  *    treatment of static objects.
  *