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.

4 files changed, 40 insertions, 14 deletions

diff --git a/includes/RtsAPI.h b/includes/RtsAPI.h
index 4748060dee..16b848678d 100644
--- a/includes/RtsAPI.h
+++ b/includes/RtsAPI.h
@@ -172,6 +172,15 @@ void rts_unlock (Capability *token);
 // when there is no current capability.
 Capability *rts_unsafeGetMyCapability (void);
 
+// Specify the Capability that the current OS thread should run on when it calls
+// into Haskell.  The actual capability will be calculated as the supplied
+// value modulo the number of enabled Capabilities.
+//
+// Note that the thread may still be migrated by the RTS scheduler, but that
+// will only happen if there are multiple threads running on one Capability and
+// another Capability is free.
+void setInCallCapability (int preferred_capability);
+
 /* ----------------------------------------------------------------------------
    Building Haskell objects from C datatypes.
    ------------------------------------------------------------------------- */
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;
             }
         }
 
diff --git a/rts/Task.c b/rts/Task.c
index 82f7780654..c30bcf17d5 100644
--- a/rts/Task.c
+++ b/rts/Task.c
@@ -213,6 +213,7 @@ newTask (rtsBool worker)
     task->n_spare_incalls = 0;
     task->spare_incalls = NULL;
     task->incall        = NULL;
+    task->preferred_capability = -1;
 
 #if defined(THREADED_RTS)
     initCondition(&task->cond);
@@ -488,6 +489,14 @@ interruptWorkerTask (Task *task)
 
 #endif /* THREADED_RTS */
 
+void
+setInCallCapability (int preferred_capability)
+{
+    Task *task = allocTask();
+    task->preferred_capability = preferred_capability;
+}
+
+
 #ifdef DEBUG
 
 void printAllTasks(void);
diff --git a/rts/Task.h b/rts/Task.h
index 37832a39d3..bcf456d270 100644
--- a/rts/Task.h
+++ b/rts/Task.h
@@ -151,6 +151,9 @@ typedef struct Task_ {
     // So that we can detect when a finalizer illegally calls back into Haskell
     rtsBool running_finalizers;
 
+    // if >= 0, this Capability will be used for in-calls
+    int preferred_capability;
+
     // Links tasks on the returning_tasks queue of a Capability, and
     // on spare_workers.
     struct Task_ *next;