Add GHC.PerCapability

8 files changed, 133 insertions, 0 deletions

diff --git a/libraries/base/GHC/PerCapability.hs b/libraries/base/GHC/PerCapability.hs
new file mode 100644
index 0000000000..e6f510b527
--- /dev/null
+++ b/libraries/base/GHC/PerCapability.hs
@@ -0,0 +1,124 @@
+{-# LANGUAGE NoImplicitPrelude #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE GADTs #-}
+
+module GHC.PerCapability
+    ( PerCapability
+    , newPerCapability
+    , getPerCapability
+    , freePerCapability
+      -- Internal
+    , capabilitiesChanged
+    ) where
+
+import GHC.Base
+import GHC.Conc.Sync (getNumCapabilities, myThreadId, threadCapability, yield)
+import GHC.Num ((-), (+))
+import Data.Foldable (mapM_, forM_)
+import Data.IORef (IORef, newIORef, readIORef, writeIORef, atomicModifyIORef)
+import GHC.IOArray (IOArray, newIOArray, readIOArray, writeIOArray,
+                    boundsIOArray)
+import GHC.IO (unsafePerformIO)
+
+-- | An array of values, one per capability
+data PerCapability a
+    = PerCapability { pcNewCap :: Int -> IO a
+                    , pcFreeCap :: a -> IO ()
+                    , pcArr :: !(IORef (IOArray Int a))
+                    }
+
+data SomePerCapability where
+    SomePerCapability :: forall a. PerCapability a -> SomePerCapability
+
+perCapabilityThings :: IORef [SomePerCapability]
+perCapabilityThings = unsafePerformIO $ newIORef []
+{-# NOINLINE perCapabilityThings #-}
+
+newPerCapability
+    :: (Int -> IO a)
+    -> (a -> IO ())
+    -> IO (PerCapability a)
+newPerCapability newCap freeCap = do
+    num_caps <- getNumCapabilities
+    arr <- newIOArray (0, num_caps-1) uninitPerCap
+    let (low, high) = boundsIOArray arr
+    forM_ [low..high] $ \n -> do
+        x <- newCap n
+        writeIOArray arr n x
+    arr_ref <- newIORef arr
+    let pc = PerCapability
+            { pcNewCap = newCap
+            , pcFreeCap = freeCap
+            , pcArr = arr_ref
+            }
+    atomicModifyIORef perCapabilityThings $ \pcs -> (SomePerCapability pc : pcs, ())
+    return pc
+
+getPerCapability
+    :: PerCapability a
+    -> IO a
+getPerCapability pc = do
+    t <- myThreadId
+    (cap, _) <- threadCapability t
+    arr <- readIORef (pcArr pc)
+    -- It is possible that we've just increased the number of capabilities and the
+    -- new EventManager has not yet been constructed by
+    -- 'ioManagerCapabilitiesChanged'. We expect this to happen very rarely.
+    -- T21561 exercises this.
+    -- Two options to proceed:
+    --  1) return the EventManager for capability 0. This is guaranteed to exist,
+    --     and "shouldn't" cause any correctness issues.
+    --  2) Busy wait, with or without a call to 'yield'. This can't deadlock,
+    --     because we must be on a brand capability and there must be a call to
+    --     'ioManagerCapabilitiesChanged' pending.
+    --
+    -- We take the second option, with the yield, judging it the most robust.
+    if inRange (boundsIOArray arr) cap
+      then readIOArray arr cap
+      else yield >> getPerCapability pc
+
+-- | Free all per-capability resources. 'getPerCapability'
+-- must not be called concurrently.
+freePerCapability
+    :: PerCapability a
+    -> IO ()
+freePerCapability pc = do
+    let err = error "freePerCapability"
+    arr <- atomicModifyIORef (pcArr pc) (\arr -> (arr, err))
+    let low, high :: Int
+        (low, high) = boundsIOArray arr
+    forM_ [low..high] $ \n -> do
+        x <- readIOArray arr n
+        pcFreeCap pc x
+
+capabilitiesChanged :: IO ()
+capabilitiesChanged = do
+    things <- readIORef perCapabilityThings
+    new_n_caps <- getNumCapabilities
+    mapM_ (change_one new_n_caps) things
+  where
+    change_one new_n_caps (SomePerCapability pc) = do
+        arr <- readIORef (pcArr pc)
+        new_arr <- newIOArray (0, new_n_caps-1) uninitPerCap
+        let (low, high) = boundsIOArray arr
+        -- Copy the existing values into the new array
+        forM_ [low .. high] $ \i -> do
+            x <- readIOArray arr i
+            writeIOArray new_arr i x
+
+        -- Free any resoures associated with dead caps
+        forM_ [new_n_caps .. high-1] $ \i -> do
+            x <- readIOArray arr i
+            pcFreeCap pc x
+
+        -- Create new resources for the new caps
+        forM_ [high+1 .. new_n_caps-1] $ \i -> do
+            x <- pcNewCap pc i
+            writeIOArray new_arr i x
+
+        -- update the array reference
+        writeIORef (pcArr pc) new_arr
+
+uninitPerCap :: a
+uninitPerCap = error "Uninitialized PerCapability slot"
+
diff --git a/libraries/base/base.cabal b/libraries/base/base.cabal
index a8875f9d59..ed0f7da5a5 100644
--- a/libraries/base/base.cabal
+++ b/libraries/base/base.cabal
@@ -350,6 +350,7 @@ Library
         GHC.Event.IntVar
         GHC.Event.PSQ
         GHC.Event.Unique
+        GHC.PerCapability
         -- GHC.IOPort -- TODO: hide again after debug
         GHC.Unicode.Internal.Bits
         GHC.Unicode.Internal.Char.DerivedCoreProperties
diff --git a/rts/Prelude.h b/rts/Prelude.h
index a474771f5d..b33b634748 100644
--- a/rts/Prelude.h
+++ b/rts/Prelude.h
@@ -61,6 +61,7 @@ PRELUDE_CLOSURE(base_GHCziConcziIO_ensureIOManagerIsRunning_closure);
 PRELUDE_CLOSURE(base_GHCziConcziIO_interruptIOManager_closure);
 PRELUDE_CLOSURE(base_GHCziConcziIO_ioManagerCapabilitiesChanged_closure);
 PRELUDE_CLOSURE(base_GHCziConcziSignal_runHandlersPtr_closure);
+PRELUDE_CLOSURE(base_GHCziPerCapability_capabilitiesChanged_closure);
 #if defined(mingw32_HOST_OS)
 PRELUDE_CLOSURE(base_GHCziEventziWindows_processRemoteCompletion_closure);
 #endif
diff --git a/rts/RtsStartup.c b/rts/RtsStartup.c
index a80dfc2959..ce4ca8a806 100644
--- a/rts/RtsStartup.c
+++ b/rts/RtsStartup.c
@@ -210,6 +210,7 @@ static void initBuiltinGcRoots(void)
     getStablePtr((StgPtr)ensureIOManagerIsRunning_closure);
     getStablePtr((StgPtr)interruptIOManager_closure);
     getStablePtr((StgPtr)ioManagerCapabilitiesChanged_closure);
+    getStablePtr((StgPtr)capabilitiesChanged_closure);
 #if !defined(mingw32_HOST_OS)
     getStablePtr((StgPtr)blockedOnBadFD_closure);
     getStablePtr((StgPtr)runHandlersPtr_closure);
diff --git a/rts/Schedule.c b/rts/Schedule.c
index 76ce2b7567..e8c6a39013 100644
--- a/rts/Schedule.c
+++ b/rts/Schedule.c
@@ -2330,6 +2330,7 @@ setNumCapabilities (uint32_t new_n_capabilities USED_IF_THREADS)
 
     // Notify IO manager that the number of capabilities has changed.
     rts_evalIO(&cap, ioManagerCapabilitiesChanged_closure, NULL);
+    rts_evalIO(&cap, capabilitiesChanged_closure, NULL);
 
     startTimer();
 
diff --git a/rts/package.conf.in b/rts/package.conf.in
index 42359f301a..74cbd0b8e4 100644
--- a/rts/package.conf.in
+++ b/rts/package.conf.in
@@ -110,6 +110,7 @@ ld-options:
          , "-Wl,-u,_base_GHCziConcziIO_interruptIOManager_closure"
          , "-Wl,-u,_base_GHCziConcziIO_ioManagerCapabilitiesChanged_closure"
          , "-Wl,-u,_base_GHCziConcziSignal_runHandlersPtr_closure"
+         , "-Wl,-u,_base_GHCziPerCapability_capabilitiesChanged_closure"
 #if defined(mingw32_HOST_OS)
          , "-Wl,-u,_base_GHCziEventziWindows_processRemoteCompletion_closure"
 #endif
@@ -223,6 +224,7 @@ ld-options:
          , "-Wl,-u,base_GHCziConcziIO_interruptIOManager_closure"
          , "-Wl,-u,base_GHCziConcziIO_ioManagerCapabilitiesChanged_closure"
          , "-Wl,-u,base_GHCziConcziSignal_runHandlersPtr_closure"
+         , "-Wl,-u,base_GHCziPerCapability_capabilitiesChanged_closure"
 #if defined(mingw32_HOST_OS)
          , "-Wl,-u,base_GHCziEventziWindows_processRemoteCompletion_closure"
 #endif
diff --git a/rts/rts.cabal.in b/rts/rts.cabal.in
index 2a113e8867..27710d6ca1 100644
--- a/rts/rts.cabal.in
+++ b/rts/rts.cabal.in
@@ -280,6 +280,7 @@ library
          "-Wl,-u,_base_GHCziConcziIO_interruptIOManager_closure"
          "-Wl,-u,_base_GHCziConcziIO_ioManagerCapabilitiesChanged_closure"
          "-Wl,-u,_base_GHCziConcziSignal_runHandlersPtr_closure"
+         "-Wl,-u,_base_GHCziPerCapability_capabilitiesChanged_closure"
          "-Wl,-u,_base_GHCziTopHandler_flushStdHandles_closure"
          "-Wl,-u,_base_GHCziTopHandler_runMainIO_closure"
          "-Wl,-u,_ghczmprim_GHCziTypes_Czh_con_info"
@@ -363,6 +364,7 @@ library
          "-Wl,-u,base_GHCziConcziIO_interruptIOManager_closure"
          "-Wl,-u,base_GHCziConcziIO_ioManagerCapabilitiesChanged_closure"
          "-Wl,-u,base_GHCziConcziSignal_runHandlersPtr_closure"
+         "-Wl,-u,base_GHCziPerCapability_capabilitiesChanged_closure"
          "-Wl,-u,base_GHCziTopHandler_flushStdHandles_closure"
          "-Wl,-u,base_GHCziTopHandler_runMainIO_closure"
          "-Wl,-u,ghczmprim_GHCziTypes_Czh_con_info"
diff --git a/rts/win32/libHSbase.def b/rts/win32/libHSbase.def
index e91d11a688..8cf184e428 100644
--- a/rts/win32/libHSbase.def
+++ b/rts/win32/libHSbase.def
@@ -29,6 +29,7 @@ EXPORTS
 	base_GHCziConcziIO_ensureIOManagerIsRunning_closure
 	base_GHCziConcziIO_interruptIOManager_closure
 	base_GHCziConcziIO_ioManagerCapabilitiesChanged_closure
+	base_GHCziPerCapability_capabilitiesChanged_closure
 	base_GHCziConcziSync_runSparks_closure
 	base_GHCziEventziWindows_processRemoteCompletion_closure