RTS: Add setInCallCapability()

This allows an OS thread to specify which capability it should run on
when it makes a call into Haskell.  It is intended for a fairly
specialised use case, when the client wants to have tighter control over
the mapping between OS threads and Capabilities - perhaps 1:1
correspondence, for example.

1 files changed, 19 insertions, 14 deletions

diff --git a/rts/Capability.c b/rts/Capability.c
index a2078e5a84..355f36d0c5 100644
--- a/rts/Capability.c
+++ b/rts/Capability.c
@@ -709,21 +709,26 @@ void waitForCapability (Capability **pCap, Task *task)
     Capability *cap = *pCap;
 
     if (cap == NULL) {
-        // Try last_free_capability first
-        cap = last_free_capability;
-        if (cap->running_task) {
-            nat i;
-            // otherwise, search for a free capability
-            cap = NULL;
-            for (i = 0; i < n_capabilities; i++) {
-                if (!capabilities[i]->running_task) {
-                    cap = capabilities[i];
-                    break;
+        if (task->preferred_capability != -1) {
+            cap = capabilities[task->preferred_capability %
+                               enabled_capabilities];
+        } else {
+            // Try last_free_capability first
+            cap = last_free_capability;
+            if (cap->running_task) {
+                nat i;
+                // otherwise, search for a free capability
+                cap = NULL;
+                for (i = 0; i < n_capabilities; i++) {
+                    if (!capabilities[i]->running_task) {
+                        cap = capabilities[i];
+                        break;
+                    }
+                }
+                if (cap == NULL) {
+                    // Can't find a free one, use last_free_capability.
+                    cap = last_free_capability;
                 }
-            }
-            if (cap == NULL) {
-                // Can't find a free one, use last_free_capability.
-                cap = last_free_capability;
             }
         }