{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnliftedFFITypes #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE Strict #-}
module Control.Concurrent.Thread.Storage
(
ThreadStorageMap
, newThreadStorageMap
, lookup
, lookupOnThread
, update
, updateOnThread
, attach
, attachOnThread
, detach
, detachFromThread
, adjust
, adjustOnThread
, storedItems
, getThreadId
#if MIN_VERSION_base(4,18,0)
, purgeDeadThreads
#endif
) where
import Control.Concurrent
import Control.Concurrent.Thread.Finalizers
import Control.Monad ( when, void )
import Control.Monad.IO.Class
import Data.Maybe (isNothing, isJust)
import GHC.Base (Addr#)
import GHC.IO (IO(..), mask_)
import GHC.Int
#if MIN_VERSION_base(4,18,0)
import GHC.Conc (listThreads)
#endif
import GHC.Conc.Sync ( ThreadId(..) )
import GHC.Prim
import qualified Data.IntMap.Strict as I
import qualified Data.IntSet as IS
import Foreign.C.Types
import Prelude hiding (lookup)
import Unsafe.Coerce (unsafeCoerce#)
foreign import ccall unsafe "rts_getThreadId" c_getThreadId :: Addr# -> CULLong
numStripes :: Int
numStripes :: Int
numStripes = Int
32
getThreadId :: ThreadId -> Int
getThreadId :: ThreadId -> Int
getThreadId (ThreadId ThreadId#
tid#) = forall a b. (Integral a, Num b) => a -> b
fromIntegral (Addr# -> CULLong
c_getThreadId (unsafeCoerce# :: forall a b. a -> b
unsafeCoerce# ThreadId#
tid#))
stripeHash :: Int -> Int
stripeHash :: Int -> Int
stripeHash = (forall a. Integral a => a -> a -> a
`mod` Int
numStripes) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Num a => a -> a
abs
readStripe :: ThreadStorageMap a -> ThreadId -> IO (I.IntMap a)
readStripe :: forall a. ThreadStorageMap a -> ThreadId -> IO (IntMap a)
readStripe (ThreadStorageMap MutableArray# RealWorld (IntMap a)
arr#) ThreadId
t = forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> forall d a.
MutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readArray# MutableArray# RealWorld (IntMap a)
arr# Int#
tid# State# RealWorld
s
where
(I# Int#
tid#) = Int -> Int
stripeHash forall a b. (a -> b) -> a -> b
$ ThreadId -> Int
getThreadId ThreadId
t
atomicModifyStripe :: ThreadStorageMap a -> Int -> (I.IntMap a -> (I.IntMap a, b)) -> IO b
atomicModifyStripe :: forall a b.
ThreadStorageMap a -> Int -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe (ThreadStorageMap MutableArray# RealWorld (IntMap a)
arr#) Int
tid IntMap a -> (IntMap a, b)
f = forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> State# RealWorld -> (# State# RealWorld, b #)
go State# RealWorld
s
where
(I# Int#
stripe#) = Int -> Int
stripeHash Int
tid
go :: State# RealWorld -> (# State# RealWorld, b #)
go State# RealWorld
s = case forall d a.
MutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readArray# MutableArray# RealWorld (IntMap a)
arr# Int#
stripe# State# RealWorld
s of
(# State# RealWorld
s1, IntMap a
intMap #) ->
let (IntMap a
updatedIntMap, b
result) = IntMap a -> (IntMap a, b)
f IntMap a
intMap
in case forall d a.
MutableArray# d a
-> Int# -> a -> a -> State# d -> (# State# d, Int#, a #)
casArray# MutableArray# RealWorld (IntMap a)
arr# Int#
stripe# IntMap a
intMap IntMap a
updatedIntMap State# RealWorld
s1 of
(# State# RealWorld
s2, Int#
outcome, IntMap a
old #) -> case Int#
outcome of
Int#
0# -> (# State# RealWorld
s2, b
result #)
Int#
1# -> State# RealWorld -> (# State# RealWorld, b #)
go State# RealWorld
s2
Int#
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Got impossible result in atomicModifyStripe"
data ThreadStorageMap a = ThreadStorageMap
(MutableArray# RealWorld (I.IntMap a))
newThreadStorageMap
:: MonadIO m => m (ThreadStorageMap a)
newThreadStorageMap :: forall (m :: * -> *) a. MonadIO m => m (ThreadStorageMap a)
newThreadStorageMap = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case forall a d.
Int# -> a -> State# d -> (# State# d, MutableArray# d a #)
newArray# Int#
numStripes# forall a. Monoid a => a
mempty State# RealWorld
s of
(# State# RealWorld
s1, MutableArray# RealWorld (IntMap a)
ma #) -> (# State# RealWorld
s1, forall a. MutableArray# RealWorld (IntMap a) -> ThreadStorageMap a
ThreadStorageMap MutableArray# RealWorld (IntMap a)
ma #)
where
(I# Int#
numStripes#) = Int
numStripes
lookup :: MonadIO m => ThreadStorageMap a -> m (Maybe a)
lookup :: forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> m (Maybe a)
lookup ThreadStorageMap a
tsm = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
ThreadId
tid <- IO ThreadId
myThreadId
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> m (Maybe a)
lookupOnThread ThreadStorageMap a
tsm ThreadId
tid
lookupOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> m (Maybe a)
lookupOnThread :: forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> m (Maybe a)
lookupOnThread ThreadStorageMap a
tsm ThreadId
tid = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
IntMap a
m <- forall a. ThreadStorageMap a -> ThreadId -> IO (IntMap a)
readStripe ThreadStorageMap a
tsm ThreadId
tid
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Int -> IntMap a -> Maybe a
I.lookup Int
threadAsInt IntMap a
m
where
threadAsInt :: Int
threadAsInt = ThreadId -> Int
getThreadId ThreadId
tid
attach :: MonadIO m => ThreadStorageMap a -> a -> m (Maybe a)
attach :: forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> a -> m (Maybe a)
attach ThreadStorageMap a
tsm a
x = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
ThreadId
tid <- IO ThreadId
myThreadId
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> a -> m (Maybe a)
attachOnThread ThreadStorageMap a
tsm ThreadId
tid a
x
attachOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> a -> m (Maybe a)
attachOnThread :: forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> a -> m (Maybe a)
attachOnThread ThreadStorageMap a
tsm ThreadId
tid a
ctxt =
forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> (Maybe a -> (Maybe a, b)) -> m b
updateOnThread ThreadStorageMap a
tsm ThreadId
tid (\Maybe a
prev -> (forall a. a -> Maybe a
Just a
ctxt, Maybe a
prev))
detach :: MonadIO m => ThreadStorageMap a -> m (Maybe a)
detach :: forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> m (Maybe a)
detach ThreadStorageMap a
tsm = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
ThreadId
tid <- IO ThreadId
myThreadId
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> m (Maybe a)
detachFromThread ThreadStorageMap a
tsm ThreadId
tid
detachFromThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> m (Maybe a)
detachFromThread :: forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> m (Maybe a)
detachFromThread ThreadStorageMap a
tsm ThreadId
tid = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
let threadAsInt :: Int
threadAsInt = ThreadId -> Int
getThreadId ThreadId
tid
forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> (Maybe a -> (Maybe a, b)) -> m b
updateOnThread ThreadStorageMap a
tsm ThreadId
tid (\Maybe a
prev -> (forall a. Maybe a
Nothing, Maybe a
prev))
updateOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> (Maybe a -> (Maybe a, b)) -> m b
updateOnThread :: forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> (Maybe a -> (Maybe a, b)) -> m b
updateOnThread ThreadStorageMap a
tsm ThreadId
tid Maybe a -> (Maybe a, b)
f = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
(Bool
isNewThreadEntry, b
result) <- forall a b.
ThreadStorageMap a -> Int -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe ThreadStorageMap a
tsm Int
threadAsInt forall a b. (a -> b) -> a -> b
$ \IntMap a
m ->
let ((Bool, b)
resultWithNewThreadDetection, IntMap a
m') =
forall (f :: * -> *) a.
Functor f =>
(Maybe a -> f (Maybe a)) -> Int -> IntMap a -> f (IntMap a)
I.alterF
(\Maybe a
x -> case Maybe a -> (Maybe a, b)
f Maybe a
x of
(Maybe a
x', b
y) -> ((forall a. Maybe a -> Bool
isNothing Maybe a
x Bool -> Bool -> Bool
&& forall a. Maybe a -> Bool
isJust Maybe a
x', b
y), Maybe a
x')
)
Int
threadAsInt
IntMap a
m
in (IntMap a
m', (Bool, b)
resultWithNewThreadDetection)
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
isNewThreadEntry forall a b. (a -> b) -> a -> b
$ do
ThreadId -> IO () -> IO ()
addThreadFinalizer ThreadId
tid forall a b. (a -> b) -> a -> b
$ forall a. ThreadStorageMap a -> Int -> IO ()
cleanUp ThreadStorageMap a
tsm Int
threadAsInt
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
result
where
threadAsInt :: Int
threadAsInt = ThreadId -> Int
getThreadId ThreadId
tid
update :: MonadIO m => ThreadStorageMap a -> (Maybe a -> (Maybe a, b)) -> m b
update :: forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> (Maybe a -> (Maybe a, b)) -> m b
update ThreadStorageMap a
tsm Maybe a -> (Maybe a, b)
f = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
ThreadId
tid <- IO ThreadId
myThreadId
forall (m :: * -> *) a b.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> (Maybe a -> (Maybe a, b)) -> m b
updateOnThread ThreadStorageMap a
tsm ThreadId
tid Maybe a -> (Maybe a, b)
f
adjust :: MonadIO m => ThreadStorageMap a -> (a -> a) -> m ()
adjust :: forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> (a -> a) -> m ()
adjust ThreadStorageMap a
tsm a -> a
f = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
ThreadId
tid <- IO ThreadId
myThreadId
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> (a -> a) -> m ()
adjustOnThread ThreadStorageMap a
tsm ThreadId
tid a -> a
f
adjustOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> (a -> a) -> m ()
adjustOnThread :: forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> (a -> a) -> m ()
adjustOnThread ThreadStorageMap a
tsm ThreadId
tid a -> a
f = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
forall a b.
ThreadStorageMap a -> Int -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe ThreadStorageMap a
tsm Int
threadAsInt forall a b. (a -> b) -> a -> b
$ \IntMap a
m -> (forall a. (a -> a) -> Int -> IntMap a -> IntMap a
I.adjust a -> a
f Int
threadAsInt IntMap a
m, ())
where
threadAsInt :: Int
threadAsInt = ThreadId -> Int
getThreadId ThreadId
tid
cleanUp :: ThreadStorageMap a -> Int -> IO ()
cleanUp :: forall a. ThreadStorageMap a -> Int -> IO ()
cleanUp ThreadStorageMap a
tsm Int
tid = do
forall a b.
ThreadStorageMap a -> Int -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe ThreadStorageMap a
tsm Int
tid forall a b. (a -> b) -> a -> b
$ \IntMap a
m ->
(forall a. Int -> IntMap a -> IntMap a
I.delete Int
tid IntMap a
m, ())
storedItems :: ThreadStorageMap a -> IO [(Int, a)]
storedItems :: forall a. ThreadStorageMap a -> IO [(Int, a)]
storedItems ThreadStorageMap a
tsm = do
[IntMap a]
stripes <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall a. ThreadStorageMap a -> Int -> IO (IntMap a)
stripeByIndex ThreadStorageMap a
tsm) [Int
0..(Int
numStripes forall a. Num a => a -> a -> a
- Int
1)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. IntMap a -> [(Int, a)]
I.toList [IntMap a]
stripes
where
stripeByIndex :: ThreadStorageMap a -> Int -> IO (I.IntMap a)
stripeByIndex :: forall a. ThreadStorageMap a -> Int -> IO (IntMap a)
stripeByIndex (ThreadStorageMap MutableArray# RealWorld (IntMap a)
arr#) (I# Int#
i#) = forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> forall d a.
MutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readArray# MutableArray# RealWorld (IntMap a)
arr# Int#
i# State# RealWorld
s
#if MIN_VERSION_base(4,18,0)
purgeDeadThreads :: MonadIO m => ThreadStorageMap a -> m ()
purgeDeadThreads tsm = do
tids <- listThreads
let threadSet = IS.fromList $ map getThreadId tids
forM_ [0..(numStripes - 1)] $ \stripe ->
atomicModifyStripe tsm stripe $ \im -> (I.restrictKeys im threadSet, ())
#endif