Allow the number of capabilities to be increased at runtime (#3729)

At present the number of capabilities can only be *increased*, not
decreased.  The latter presents a few more challenges!

1 files changed, 48 insertions, 18 deletions

diff --git a/rts/Capability.c b/rts/Capability.c
index 4d23f71a86..7ce23a12e6 100644
--- a/rts/Capability.c
+++ b/rts/Capability.c
@@ -27,6 +27,8 @@
 #include "STM.h"
 #include "RtsUtils.h"
 
+#include <string.h>
+
 // one global capability, this is the Capability for non-threaded
 // builds, and for +RTS -N1
 Capability MainCapability;
@@ -299,7 +301,6 @@ initCapabilities( void )
     traceCapsetCreate(CAPSET_CLOCKDOMAIN_DEFAULT, CapsetTypeClockdomain);
 
 #if defined(THREADED_RTS)
-    nat i;
 
 #ifndef REG_Base
     // We can't support multiple CPUs if BaseReg is not a register
@@ -309,24 +310,10 @@ initCapabilities( void )
     }
 #endif
 
+    n_capabilities = 0;
+    moreCapabilities(0, RtsFlags.ParFlags.nNodes);
     n_capabilities = RtsFlags.ParFlags.nNodes;
 
-    if (n_capabilities == 1) {
-	capabilities = &MainCapability;
-	// THREADED_RTS must work on builds that don't have a mutable
-	// BaseReg (eg. unregisterised), so in this case
-	// capabilities[0] must coincide with &MainCapability.
-    } else {
-	capabilities = stgMallocBytes(n_capabilities * sizeof(Capability),
-				      "initCapabilities");
-    }
-
-    for (i = 0; i < n_capabilities; i++) {
-	initCapability(&capabilities[i], i);
-    }
-
-    debugTrace(DEBUG_sched, "allocated %d capabilities", n_capabilities);
-
 #else /* !THREADED_RTS */
 
     n_capabilities = 1;
@@ -341,6 +328,46 @@ initCapabilities( void )
     last_free_capability = &capabilities[0];
 }
 
+Capability *
+moreCapabilities (nat from USED_IF_THREADS, nat to USED_IF_THREADS)
+{
+#if defined(THREADED_RTS)
+    nat i;
+    Capability *old_capabilities = capabilities;
+
+    if (to == 1) {
+        // THREADED_RTS must work on builds that don't have a mutable
+        // BaseReg (eg. unregisterised), so in this case
+	// capabilities[0] must coincide with &MainCapability.
+        capabilities = &MainCapability;
+    } else {
+        capabilities = stgMallocBytes(to * sizeof(Capability),
+                                      "moreCapabilities");
+
+        if (from > 0) {
+            memcpy(capabilities, old_capabilities, from * sizeof(Capability));
+        }
+    }
+
+    for (i = from; i < to; i++) {
+	initCapability(&capabilities[i], i);
+    }
+
+    last_free_capability = NULL;
+
+    debugTrace(DEBUG_sched, "allocated %d more capabilities", to - from);
+
+    // Return the old array to free later.
+    if (from > 1) {
+        return old_capabilities;
+    } else {
+        return NULL;
+    }
+#else
+    return NULL;
+#endif
+}
+
 /* ----------------------------------------------------------------------------
  * setContextSwitches: cause all capabilities to context switch as
  * soon as possible.
@@ -426,7 +453,10 @@ releaseCapability_ (Capability* cap,
 	return;
     }
 
-    if (pending_sync == SYNC_GC_SEQ || pending_sync == SYNC_FORK) {
+    // If there is a pending sync, then we should just leave the
+    // Capability free.  The thread trying to sync will be about to
+    // call waitForReturnCapability().
+    if (pending_sync != 0 && pending_sync != SYNC_GC_PAR) {
       last_free_capability = cap; // needed?
       debugTrace(DEBUG_sched, "sync pending, set capability %d free", cap->no);
       return;