rts: forkOn context switches the target capability

Fixes #21824

4 files changed, 14 insertions, 8 deletions

diff --git a/rts/Capability.c b/rts/Capability.c
index abb3e5ed8c..c3571f6d64 100644
--- a/rts/Capability.c
+++ b/rts/Capability.c
@@ -466,7 +466,7 @@ void contextSwitchAllCapabilities(void)
 {
     uint32_t i;
     for (i=0; i < n_capabilities; i++) {
-        contextSwitchCapability(capabilities[i]);
+        contextSwitchCapability(capabilities[i], true);
     }
 }
 
diff --git a/rts/Capability.h b/rts/Capability.h
index 4a27e618fe..1f3ff95736 100644
--- a/rts/Capability.h
+++ b/rts/Capability.h
@@ -365,7 +365,10 @@ void shutdownCapabilities(Task *task, bool wait_foreign);
 
 // cause all capabilities to context switch as soon as possible.
 void contextSwitchAllCapabilities(void);
-INLINE_HEADER void contextSwitchCapability(Capability *cap);
+
+// if immediately is set then the capability will context-switch at the next
+// heap-check.  Otherwise it will context switch at the next failing heap-check.
+INLINE_HEADER void contextSwitchCapability(Capability *cap, bool immediately);
 
 // cause all capabilities to stop running Haskell code and return to
 // the scheduler as soon as possible.
@@ -478,9 +481,11 @@ interruptCapability (Capability *cap)
 }
 
 INLINE_HEADER void
-contextSwitchCapability (Capability *cap)
+contextSwitchCapability (Capability *cap, bool immediately)
 {
-    stopCapability(cap);
+    if(immediately) {
+        stopCapability(cap);
+    }
     SEQ_CST_STORE(&cap->context_switch, true);
 }
 
diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index 430824e155..7b760e5702 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -1062,10 +1062,6 @@ again: MAYBE_GC(again);
 
     ccall scheduleThreadOn(MyCapability() "ptr", cpu, threadid "ptr");
 
-    // context switch soon, but not immediately: we don't want every
-    // forkIO to force a context-switch.
-    Capability_context_switch(MyCapability()) = 1 :: CInt;
-
     return (threadid);
 }
 
diff --git a/rts/Schedule.c b/rts/Schedule.c
index 0a75bd1a55..fa0fc8c63e 100644
--- a/rts/Schedule.c
+++ b/rts/Schedule.c
@@ -2586,6 +2586,9 @@ scheduleThreadOn(Capability *cap, StgWord cpu USED_IF_THREADS, StgTSO *tso)
 {
     tso->flags |= TSO_LOCKED; // we requested explicit affinity; don't
                               // move this thread from now on.
+
+    // We will context switch soon, but not immediately: we don't want every
+    // fork to force a context-switch.
 #if defined(THREADED_RTS)
     cpu %= enabled_capabilities;
     if (cpu == cap->no) {
@@ -2593,8 +2596,10 @@ scheduleThreadOn(Capability *cap, StgWord cpu USED_IF_THREADS, StgTSO *tso)
     } else {
         migrateThread(cap, tso, capabilities[cpu]);
     }
+    contextSwitchCapability(capabilities[cpu], false);
 #else
     appendToRunQueue(cap,tso);
+    contextSwitchCapability(cap, false);
 #endif
 }