Fix hs_try_putmvar losing track of running cap

If hs_try_putmvar was called through an unsafe import, it would lose track of the running cap causing a deadlock

3 files changed, 26 insertions, 0 deletions

diff --git a/rts/RtsAPI.c b/rts/RtsAPI.c
index 26433ac209..51a1f2b7cf 100644
--- a/rts/RtsAPI.c
+++ b/rts/RtsAPI.c
@@ -682,6 +682,7 @@ void hs_try_putmvar (/* in */ int capability,
 {
     Task *task = getTask();
     Capability *cap;
+    Capability *task_old_cap USED_IF_THREADS;
 
     if (capability < 0) {
         capability = task->preferred_capability;
@@ -702,6 +703,7 @@ void hs_try_putmvar (/* in */ int capability,
     // If the capability is free, we can perform the tryPutMVar immediately
     if (cap->running_task == NULL) {
         cap->running_task = task;
+        task_old_cap = task->cap;
         task->cap = cap;
         RELEASE_LOCK(&cap->lock);
 
@@ -712,6 +714,7 @@ void hs_try_putmvar (/* in */ int capability,
         // Wake up the capability, which will start running the thread that we
         // just awoke (if there was one).
         releaseCapability(cap);
+        task->cap = task_old_cap;
     } else {
         PutMVar *p = stgMallocBytes(sizeof(PutMVar),"hs_try_putmvar");
         // We cannot deref the StablePtr if we don't have a capability,
diff --git a/testsuite/tests/rts/T15427.hs b/testsuite/tests/rts/T15427.hs
new file mode 100644
index 0000000000..d4af150106
--- /dev/null
+++ b/testsuite/tests/rts/T15427.hs
@@ -0,0 +1,21 @@
+import Control.Concurrent.MVar
+import Control.Monad
+import Foreign.C.Types (CInt(..))
+import Foreign.StablePtr (StablePtr)
+import GHC.Conc
+
+foreign import ccall unsafe hs_try_putmvar :: CInt -> StablePtr PrimMVar -> IO ()
+
+main = do
+  mvs <- forM [0..numCapabilities] (\idx -> do
+                                         a <- newEmptyMVar
+                                         b <- newEmptyMVar
+                                         return $ (idx, a, b))
+  forM_ [mvs, reverse mvs] $ \mvars -> do
+    forM_ mvars $ (\(cap,a,b) -> forkOn cap $ do
+                      takeMVar a
+                      putMVar b ())
+    forM_ mvars $ \(cap, a, _) -> do
+      sp <- newStablePtrPrimMVar a
+      hs_try_putmvar (fromIntegral cap) sp
+    forM_ mvars $ \(_,_,b) -> takeMVar b
diff --git a/testsuite/tests/rts/all.T b/testsuite/tests/rts/all.T
index e4e2561c2e..ca0e652a48 100644
--- a/testsuite/tests/rts/all.T
+++ b/testsuite/tests/rts/all.T
@@ -417,3 +417,5 @@ test('InitEventLogging',
 test('T17088',
      [only_ways(['normal']), extra_run_opts('+RTS -c -A256k -RTS')],
      compile_and_run, ['-rtsopts -O2'])
+
+test('T15427', normal, compile_and_run, [''])