{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnliftedFFITypes #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE BangPatterns #-}
-- | A perilous implementation of thread-local storage for Haskell.
-- This module uses a fair amount of GHC internals to enable performing
-- lookups of context for any threads that are alive. Caution should be
-- taken for consumers of this module to not retain ThreadId references
-- indefinitely, as that could delay cleanup of thread-local state.
--
-- Thread-local contexts have the following semantics:
--
-- - A value 'attach'ed to a 'ThreadId' will remain alive at least as long
--   as the 'ThreadId'. 
-- - A value may be detached from a 'ThreadId' via 'detach' by the
--   library consumer without detriment.
-- - No guarantees are made about when a value will be garbage-collected
--   once all references to 'ThreadId' have been dropped. However, this simply
--   means in practice that any unused contexts will cleaned up upon the next
--   garbage collection and may not be actively freed when the program exits.
--
-- Note that this implementation of context sharing is
-- mildly expensive for the garbage collector, hard to reason about without deep
-- knowledge of the code you are instrumenting, and has limited guarantees of behavior 
-- across GHC versions due to internals usage.
module Control.Concurrent.Thread.Storage 
  ( 
    -- * Create a 'ThreadStorageMap'
    ThreadStorageMap
  , newThreadStorageMap
    -- * Retrieve values from a 'ThreadStorageMap'
  , lookup
  , lookupOnThread
    -- * Update values in a 'ThreadStorageMap'
  , update
  , updateOnThread
    -- * Associate values with a thread in a 'ThreadStorageMap'
  , attach
  , attachOnThread
    -- * Remove values from a thread in a 'ThreadStorageMap'
  , detach
  , detachFromThread
    -- * Update values for a thread in a 'ThreadStorageMap'
  , adjust
  , adjustOnThread
    -- * Monitoring utilities
  , storedItems
    -- * Thread ID manipulation
  , getThreadId
#if MIN_VERSION_base(4,18,0)
  , purgeDeadThreads
#endif
  ) where

import Control.Concurrent
import Control.Concurrent.Thread.Finalizers
import Control.Monad ( when, void, forM_ )
import Control.Monad.IO.Class
import Data.Maybe (isNothing, isJust)
import Data.Word (Word64)
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 GHC.Exts (unsafeCoerce#)

foreign import ccall unsafe "rts_getThreadId" c_getThreadId :: Addr# -> CULLong

numStripes :: Word
numStripes :: Word
numStripes = Word
32

getThreadId :: ThreadId -> Word
getThreadId :: ThreadId -> Word
getThreadId (ThreadId ThreadId#
tid#) = CULLong -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Addr# -> CULLong
c_getThreadId (ThreadId# -> Addr#
forall a b. a -> b
unsafeCoerce# ThreadId#
tid#))

stripeHash :: Word -> Int
stripeHash :: Word -> Int
stripeHash = Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int) -> (Word -> Word) -> Word -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> Word -> Word
forall a. Integral a => a -> a -> a
`mod` Word
numStripes)

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 = (State# RealWorld -> (# State# RealWorld, IntMap a #))
-> IO (IntMap a)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, IntMap a #))
 -> IO (IntMap a))
-> (State# RealWorld -> (# State# RealWorld, IntMap a #))
-> IO (IntMap a)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> MutableArray# RealWorld (IntMap a)
-> Int# -> State# RealWorld -> (# State# RealWorld, IntMap a #)
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#) = Word -> Int
stripeHash (Word -> Int) -> Word -> Int
forall a b. (a -> b) -> a -> b
$ ThreadId -> Word
getThreadId ThreadId
t

atomicModifyStripe :: ThreadStorageMap a -> Word -> (I.IntMap a -> (I.IntMap a, b)) -> IO b
atomicModifyStripe :: forall a b.
ThreadStorageMap a -> Word -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe (ThreadStorageMap MutableArray# RealWorld (IntMap a)
arr#) Word
tid IntMap a -> (IntMap a, b)
f = (State# RealWorld -> (# State# RealWorld, b #)) -> IO b
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, b #)) -> IO b)
-> (State# RealWorld -> (# State# RealWorld, b #)) -> IO b
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
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Word -> Int
stripeHash Word
tid
    go :: State# RealWorld -> (# State# RealWorld, b #)
go State# RealWorld
s = case MutableArray# RealWorld (IntMap a)
-> Int# -> State# RealWorld -> (# State# RealWorld, IntMap a #)
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 MutableArray# RealWorld (IntMap a)
-> Int#
-> IntMap a
-> IntMap a
-> State# RealWorld
-> (# State# RealWorld, Int#, IntMap a #)
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#
_ -> [Char] -> (# State# RealWorld, b #)
forall a. HasCallStack => [Char] -> a
error [Char]
"Got impossible result in atomicModifyStripe"
          
-- | A storage mechanism for values of a type. This structure retains items
-- on per-(green)thread basis, which can be useful in rare cases.
data ThreadStorageMap a = ThreadStorageMap 
  (MutableArray# RealWorld (I.IntMap a))

-- | Create a new thread storage map. The map is striped by thread
-- into 32 sections in order to reduce contention.
newThreadStorageMap 
  :: MonadIO m => m (ThreadStorageMap a)
newThreadStorageMap :: forall (m :: * -> *) a. MonadIO m => m (ThreadStorageMap a)
newThreadStorageMap = IO (ThreadStorageMap a) -> m (ThreadStorageMap a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (ThreadStorageMap a) -> m (ThreadStorageMap a))
-> IO (ThreadStorageMap a) -> m (ThreadStorageMap a)
forall a b. (a -> b) -> a -> b
$ (State# RealWorld -> (# State# RealWorld, ThreadStorageMap a #))
-> IO (ThreadStorageMap a)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, ThreadStorageMap a #))
 -> IO (ThreadStorageMap a))
-> (State# RealWorld -> (# State# RealWorld, ThreadStorageMap a #))
-> IO (ThreadStorageMap a)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case Int#
-> IntMap a
-> State# RealWorld
-> (# State# RealWorld, MutableArray# RealWorld (IntMap a) #)
forall a d.
Int# -> a -> State# d -> (# State# d, MutableArray# d a #)
newArray# Int#
numStripes# IntMap a
forall a. Monoid a => a
mempty State# RealWorld
s of
  (# State# RealWorld
s1, MutableArray# RealWorld (IntMap a)
ma #) -> (# State# RealWorld
s1, MutableArray# RealWorld (IntMap a) -> ThreadStorageMap a
forall a. MutableArray# RealWorld (IntMap a) -> ThreadStorageMap a
ThreadStorageMap MutableArray# RealWorld (IntMap a)
ma #)
  where
    (I# Int#
numStripes#) = Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
numStripes

-- | Retrieve a value if it exists for the current thread
lookup :: MonadIO m => ThreadStorageMap a -> m (Maybe a)
lookup :: forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> m (Maybe a)
lookup ThreadStorageMap a
tsm = IO (Maybe a) -> m (Maybe a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe a) -> m (Maybe a)) -> IO (Maybe a) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ do
  ThreadId
tid <- IO ThreadId
myThreadId
  ThreadStorageMap a -> ThreadId -> IO (Maybe a)
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> m (Maybe a)
lookupOnThread ThreadStorageMap a
tsm ThreadId
tid

-- | Retrieve a value if it exists for the specified thread
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 = IO (Maybe a) -> m (Maybe a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe a) -> m (Maybe a)) -> IO (Maybe a) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ do
  IntMap a
m <- ThreadStorageMap a -> ThreadId -> IO (IntMap a)
forall a. ThreadStorageMap a -> ThreadId -> IO (IntMap a)
readStripe ThreadStorageMap a
tsm ThreadId
tid
  Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
I.lookup Int
threadAsInt IntMap a
m
  where 
    threadAsInt :: Int
threadAsInt = Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int) -> Word -> Int
forall a b. (a -> b) -> a -> b
$ ThreadId -> Word
getThreadId ThreadId
tid

-- | Associate the provided value with the current thread.
--
-- Returns the previous value if it was set.
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 = IO (Maybe a) -> m (Maybe a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe a) -> m (Maybe a)) -> IO (Maybe a) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ do
  ThreadId
tid <- IO ThreadId
myThreadId
  ThreadStorageMap a -> ThreadId -> a -> IO (Maybe a)
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> a -> m (Maybe a)
attachOnThread ThreadStorageMap a
tsm ThreadId
tid a
x

-- | Associate the provided value with the specified thread. This replaces
-- any values already associated with the 'ThreadId'.
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 = 
  ThreadStorageMap a
-> ThreadId -> (Maybe a -> (Maybe a, Maybe a)) -> m (Maybe a)
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 -> (a -> Maybe a
forall a. a -> Maybe a
Just a
ctxt, Maybe a
prev))

-- | Disassociate the associated value from the current thread, returning it if it exists.
detach :: MonadIO m => ThreadStorageMap a -> m (Maybe a)
detach :: forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> m (Maybe a)
detach ThreadStorageMap a
tsm = IO (Maybe a) -> m (Maybe a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe a) -> m (Maybe a)) -> IO (Maybe a) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ do
  ThreadId
tid <- IO ThreadId
myThreadId
  ThreadStorageMap a -> ThreadId -> IO (Maybe a)
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> m (Maybe a)
detachFromThread ThreadStorageMap a
tsm ThreadId
tid

-- | Disassociate the associated value from the specified thread, returning it if it exists.
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 = IO (Maybe a) -> m (Maybe a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe a) -> m (Maybe a)) -> IO (Maybe a) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ do
  let threadAsInt :: Word
threadAsInt = ThreadId -> Word
getThreadId ThreadId
tid
  ThreadStorageMap a
-> ThreadId -> (Maybe a -> (Maybe a, Maybe a)) -> IO (Maybe a)
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 -> (Maybe a
forall a. Maybe a
Nothing, Maybe a
prev))

-- | The most general function in this library. Update a 'ThreadStorageMap' on a given thread,
-- with the ability to add or remove values and return some sort of result.
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 = IO b -> m b
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> IO b -> m b
forall a b. (a -> b) -> a -> b
$ IO b -> IO b
forall a. IO a -> IO a
mask_ (IO b -> IO b) -> IO b -> IO b
forall a b. (a -> b) -> a -> b
$ do
  -- ^ We mask here in order to ensure that the finalizer will always be created
  (Bool
isNewThreadEntry, b
result) <- ThreadStorageMap a
-> Word -> (IntMap a -> (IntMap a, (Bool, b))) -> IO (Bool, b)
forall a b.
ThreadStorageMap a -> Word -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe ThreadStorageMap a
tsm Word
threadAsWord ((IntMap a -> (IntMap a, (Bool, b))) -> IO (Bool, b))
-> (IntMap a -> (IntMap a, (Bool, b))) -> IO (Bool, b)
forall a b. (a -> b) -> a -> b
$ \IntMap a
m -> 
    let ((Bool, b)
resultWithNewThreadDetection, IntMap a
m') = 
          (Maybe a -> ((Bool, b), Maybe a))
-> Int -> IntMap a -> ((Bool, b), IntMap a)
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) -> ((Maybe a -> Bool
forall a. Maybe a -> Bool
isNothing Maybe a
x Bool -> Bool -> Bool
&& Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
x', b
y), Maybe a
x')
            ) 
            (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
threadAsWord)
            IntMap a
m
     in (IntMap a
m', (Bool, b)
resultWithNewThreadDetection)
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
isNewThreadEntry (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    ThreadId -> IO () -> IO ()
addThreadFinalizer ThreadId
tid (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ThreadStorageMap a -> Word -> IO ()
forall a. ThreadStorageMap a -> Word -> IO ()
cleanUp ThreadStorageMap a
tsm Word
threadAsWord
  b -> IO b
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
result
  where 
    threadAsWord :: Word
threadAsWord = ThreadId -> Word
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 = IO b -> m b
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> IO b -> m b
forall a b. (a -> b) -> a -> b
$ do
  ThreadId
tid <- IO ThreadId
myThreadId
  ThreadStorageMap a -> ThreadId -> (Maybe a -> (Maybe a, b)) -> IO b
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

-- | Update the associated value for the current thread if it is attached.
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 = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
  ThreadId
tid <- IO ThreadId
myThreadId
  ThreadStorageMap a -> ThreadId -> (a -> a) -> IO ()
forall (m :: * -> *) a.
MonadIO m =>
ThreadStorageMap a -> ThreadId -> (a -> a) -> m ()
adjustOnThread ThreadStorageMap a
tsm ThreadId
tid a -> a
f

-- | Update the associated value for the specified thread if it is attached.
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 = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
  ThreadStorageMap a -> Word -> (IntMap a -> (IntMap a, ())) -> IO ()
forall a b.
ThreadStorageMap a -> Word -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe ThreadStorageMap a
tsm Word
threadAsWord ((IntMap a -> (IntMap a, ())) -> IO ())
-> (IntMap a -> (IntMap a, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \IntMap a
m -> ((a -> a) -> Int -> IntMap a -> IntMap a
forall a. (a -> a) -> Int -> IntMap a -> IntMap a
I.adjust a -> a
f (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
threadAsWord) IntMap a
m, ())
  where 
    threadAsWord :: Word
threadAsWord = ThreadId -> Word
getThreadId ThreadId
tid 

-- Remove this context for thread from the map on finalization
cleanUp :: ThreadStorageMap a -> Word -> IO ()
cleanUp :: forall a. ThreadStorageMap a -> Word -> IO ()
cleanUp ThreadStorageMap a
tsm Word
tid = do
  ThreadStorageMap a -> Word -> (IntMap a -> (IntMap a, ())) -> IO ()
forall a b.
ThreadStorageMap a -> Word -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe ThreadStorageMap a
tsm Word
tid ((IntMap a -> (IntMap a, ())) -> IO ())
-> (IntMap a -> (IntMap a, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \IntMap a
m -> 
    (Int -> IntMap a -> IntMap a
forall a. Int -> IntMap a -> IntMap a
I.delete (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
tid) IntMap a
m, ())

-- | List thread ids with live entries in the 'ThreadStorageMap'.
-- 
-- This is useful for monitoring purposes to verify that there
-- are no memory leaks retaining threads and thus preventing
-- items from being freed from a 'ThreadStorageMap' 
storedItems :: ThreadStorageMap a -> IO [(Int, a)]
storedItems :: forall a. ThreadStorageMap a -> IO [(Int, a)]
storedItems ThreadStorageMap a
tsm = do
  [IntMap a]
stripes <- (Int -> IO (IntMap a)) -> [Int] -> IO [IntMap a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (ThreadStorageMap a -> Int -> IO (IntMap a)
forall a. ThreadStorageMap a -> Int -> IO (IntMap a)
stripeByIndex ThreadStorageMap a
tsm) [Int
0..(Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
numStripes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]
  [(Int, a)] -> IO [(Int, a)]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(Int, a)] -> IO [(Int, a)]) -> [(Int, a)] -> IO [(Int, a)]
forall a b. (a -> b) -> a -> b
$ (IntMap a -> [(Int, a)]) -> [IntMap a] -> [(Int, a)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap IntMap a -> [(Int, a)]
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#) = (State# RealWorld -> (# State# RealWorld, IntMap a #))
-> IO (IntMap a)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, IntMap a #))
 -> IO (IntMap a))
-> (State# RealWorld -> (# State# RealWorld, IntMap a #))
-> IO (IntMap a)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> MutableArray# RealWorld (IntMap a)
-> Int# -> State# RealWorld -> (# State# RealWorld, IntMap a #)
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)
-- | This should generally not be needed, but may be used to remove values prior to GC-triggered finalizers being run from the 'ThreadStorageMap' for threads that have exited.
purgeDeadThreads :: MonadIO m => ThreadStorageMap a -> m ()
purgeDeadThreads :: forall (m :: * -> *) a. MonadIO m => ThreadStorageMap a -> m ()
purgeDeadThreads ThreadStorageMap a
tsm = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
  [ThreadId]
tids <- IO [ThreadId]
listThreads
  let threadSet :: IntSet
threadSet = [Int] -> IntSet
IS.fromList ([Int] -> IntSet) -> [Int] -> IntSet
forall a b. (a -> b) -> a -> b
$ (ThreadId -> Int) -> [ThreadId] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int) -> (ThreadId -> Word) -> ThreadId -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ThreadId -> Word
getThreadId) [ThreadId]
tids
  [Word] -> (Word -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Word
0..(Word
numStripes Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1)] ((Word -> IO ()) -> IO ()) -> (Word -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Word
stripe ->
    ThreadStorageMap a -> Word -> (IntMap a -> (IntMap a, ())) -> IO ()
forall a b.
ThreadStorageMap a -> Word -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe ThreadStorageMap a
tsm Word
stripe ((IntMap a -> (IntMap a, ())) -> IO ())
-> (IntMap a -> (IntMap a, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \IntMap a
im -> (IntMap a -> IntSet -> IntMap a
forall a. IntMap a -> IntSet -> IntMap a
I.restrictKeys IntMap a
im IntSet
threadSet, ())
#endif