When using rts_setInCallCapability, lock incall threads

This diff makes sure that incall threads, when using `rts_setInCallCapability`, will be created as locked.
If the thread is not locked, the thread might end up being scheduled to a different capability.
While this is mentioned in the docs for `rts_setInCallCapability,`, it makes the method significantly less useful as there is no guarantees on the capability being used.

This commit also adds a test to make sure things stay on the correct capability.

3 files changed, 116 insertions, 0 deletions

diff --git a/testsuite/tests/ffi/should_run/IncallAffinity.hs b/testsuite/tests/ffi/should_run/IncallAffinity.hs
new file mode 100644
index 0000000000..386e9950e8
--- /dev/null
+++ b/testsuite/tests/ffi/should_run/IncallAffinity.hs
@@ -0,0 +1,36 @@
+module Lib (capTest) where
+
+import Control.Concurrent
+import Control.Exception
+import Control.Concurrent.MVar
+import Control.Monad (when)
+import System.Exit
+
+foreign export ccall "capTest" capTest :: IO Int
+
+capTest :: IO Int
+capTest = catch go handle
+  where
+    handle :: SomeException -> IO Int
+    handle e = do
+      putStrLn $ "Failed " ++ (show e)
+      return (-1)
+    getCap = fmap fst $ threadCapability =<< myThreadId
+    go = do
+      when (not rtsSupportsBoundThreads) $
+        die "This test requires -threaded"
+      mvar <- newEmptyMVar
+      mvar2 <- newEmptyMVar
+      (cap, locked) <- threadCapability =<< myThreadId
+      forkOn cap $ do
+        putMVar mvar =<< getCap
+        takeMVar mvar2
+      -- if cap is locked, then this would get scheduled on a different
+      -- capacity.
+      fCap <- takeMVar mvar
+      putMVar mvar2 ()
+      cap2 <- getCap
+      when (fCap /= cap) (fail "expected cap to be the same")
+      when (cap2 /= cap) (fail "expected cap to be the same when returning")
+      when (not locked) (fail "expected to be locked")
+      return cap
diff --git a/testsuite/tests/ffi/should_run/IncallAffinity_c.c b/testsuite/tests/ffi/should_run/IncallAffinity_c.c
new file mode 100644
index 0000000000..bd719dff99
--- /dev/null
+++ b/testsuite/tests/ffi/should_run/IncallAffinity_c.c
@@ -0,0 +1,78 @@
+#include "HsFFI.h"
+
+#include <stdio.h>
+#include "Rts.h"
+#include <pthread.h>
+
+#define THREADS 6
+#define OK 9999
+static OSThreadId ids[THREADS];
+static int results[THREADS];
+static int waiters = 0;
+static int done = 0;
+static Condition cond;
+static Mutex mutex;
+
+HsInt capTest();
+
+void* OSThreadProcAttr go(void *info)
+{
+  int cap;
+  int res;
+  int threadNum = *(int*)(info);
+
+  // divide everything onto two caps (if there are two)
+  cap = (threadNum % 2) % enabled_capabilities;
+
+  OS_ACQUIRE_LOCK(&mutex);
+  waiters++;
+  if (waiters == THREADS) {
+    broadcastCondition(&cond);
+  } else {
+    while(waiters != THREADS) {
+      waitCondition(&cond, &mutex);
+    }
+  }
+  OS_RELEASE_LOCK(&mutex);
+
+  rts_setInCallCapability(cap, 0);
+  res = capTest();
+  *(int*)info = res == cap ? OK : res;
+  OS_ACQUIRE_LOCK(&mutex);
+  done++;
+  broadcastCondition(&cond);
+  OS_RELEASE_LOCK(&mutex);
+  return 0;
+}
+
+int main(int argc, char *argv[])
+{
+  int n;
+  bool ok;
+  hs_init(&argc, &argv);
+  initCondition(&cond);
+  initMutex(&mutex);
+  waiters = 0;
+  done = 0;
+  ok = true;
+  for (n=0; n < THREADS; n++) {
+    results[n] = n;
+    if (createOSThread(&ids[n], "test", go, (void*)&results[n])) {
+      printf("unable to create thread %d\n", n);
+      exit(1);
+    }
+  }
+  OS_ACQUIRE_LOCK(&mutex);
+  while(done != THREADS) {
+    waitCondition(&cond, &mutex);
+  }
+  OS_RELEASE_LOCK(&mutex);
+  for (n = 0; n < THREADS; n++) {
+    if (results[n] != OK) {
+      printf("%d: unexpected result was %d\n", n, results[n]);
+      ok = false;
+    }
+  }
+  hs_exit();
+  return ok ? 0 : 1;
+}
diff --git a/testsuite/tests/ffi/should_run/all.T b/testsuite/tests/ffi/should_run/all.T
index fb840861e6..bde21c7c26 100644
--- a/testsuite/tests/ffi/should_run/all.T
+++ b/testsuite/tests/ffi/should_run/all.T
@@ -218,3 +218,5 @@ test('UnliftedNewtypesByteArrayOffset', [omit_ways(['ghci'])], compile_and_run,
 
 test('T17471', [omit_ways(['ghci'])], compile_and_run,
      ['T17471_c.c -optc-D -optcFOO'])
+
+test('IncallAffinity', [req_smp, only_ways(['threaded1', 'threaded2'])], compile_and_run, ['IncallAffinity_c.c -no-hs-main'])