{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnliftedFFITypes #-}
{-# LANGUAGE UnboxedTuples #-}
-- | 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: This library assumes that 'ThreadId's aren't reused before a finalizer runs to
-- clean up the 'ThreadStorageMap', so quick thread churn of this nature.
--
-- Also note: This implementation of context sharing is
-- mildly expensive (~40ns to 'attach' a value) relative to using pure code / idiomatic things 
-- like MonadReader, hard to reason about without deep knowledge of threading in the code you are
-- using, 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
    -- * 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
  ) where

import Control.Concurrent
import Control.Concurrent.Thread.Finalizers
import Control.Monad ( void )
import Control.Monad.IO.Class
import Data.IORef
import GHC.IO (IO(..))
import GHC.Int
import GHC.Conc.Sync ( ThreadId(..) )
import GHC.Prim
import qualified Data.IntMap.Strict as I
import Foreign.C.Types
import Prelude hiding (lookup)

foreign import ccall unsafe "rts_getThreadId" c_getThreadId :: ThreadId# -> CInt

getThreadId :: ThreadId -> Int
getThreadId :: ThreadId -> Int
getThreadId (ThreadId ThreadId#
tid#) = CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ThreadId# -> CInt
c_getThreadId ThreadId#
tid#)
{-# INLINE getThreadId #-}
          
-- | A storage mechanism for values of a type. This structure retains items
-- on per-(green)thread basis, which can be useful in rare cases.
newtype ThreadStorageMap a = ThreadStorageMap (IORef (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 :: m (ThreadStorageMap a)
newThreadStorageMap = IO (ThreadStorageMap a) -> m (ThreadStorageMap a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IORef (IntMap a) -> ThreadStorageMap a
forall a. IORef (IntMap a) -> ThreadStorageMap a
ThreadStorageMap (IORef (IntMap a) -> ThreadStorageMap a)
-> IO (IORef (IntMap a)) -> IO (ThreadStorageMap a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntMap a -> IO (IORef (IntMap a))
forall a. a -> IO (IORef a)
newIORef IntMap a
forall a. Monoid a => a
mempty)

readStripe :: ThreadStorageMap a -> IO (I.IntMap a)
readStripe :: ThreadStorageMap a -> IO (IntMap a)
readStripe (ThreadStorageMap IORef (IntMap a)
tsm) = IORef (IntMap a) -> IO (IntMap a)
forall a. IORef a -> IO a
readIORef IORef (IntMap a)
tsm
{-# INLINABLE readStripe #-}

atomicModifyStripe :: ThreadStorageMap a -> (I.IntMap a -> (I.IntMap a, b)) -> IO b
atomicModifyStripe :: ThreadStorageMap a -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe (ThreadStorageMap IORef (IntMap a)
tsm) IntMap a -> (IntMap a, b)
f = IORef (IntMap a) -> (IntMap a -> (IntMap a, b)) -> IO b
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (IntMap a)
tsm IntMap a -> (IntMap a, b)
f
{-# INLINABLE atomicModifyStripe #-}

-- | Retrieve a value if it exists for the current thread
lookup :: MonadIO m => ThreadStorageMap a -> m (Maybe a)
lookup :: ThreadStorageMap a -> m (Maybe a)
lookup ThreadStorageMap a
tsm = IO (Maybe a) -> m (Maybe 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
{-# INLINABLE lookup #-}

-- | Retrieve a value if it exists for the specified thread
lookupOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> m (Maybe a)
lookupOnThread :: ThreadStorageMap a -> ThreadId -> m (Maybe a)
lookupOnThread ThreadStorageMap a
tsm ThreadId
tid = IO (Maybe a) -> m (Maybe 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 :: Int
threadAsInt = ThreadId -> Int
getThreadId ThreadId
tid
  IntMap a
m <- ThreadStorageMap a -> IO (IntMap a)
forall a. ThreadStorageMap a -> IO (IntMap a)
readStripe ThreadStorageMap a
tsm
  Maybe a -> IO (Maybe 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
{-# INLINABLE lookupOnThread #-}

-- | Associate the provided value with the current thread
attach :: MonadIO m => ThreadStorageMap a -> a -> m (Maybe a)
attach :: ThreadStorageMap a -> a -> m (Maybe a)
attach ThreadStorageMap a
tsm a
x = IO (Maybe a) -> m (Maybe 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
{-# INLINABLE attach #-}

-- | Associate the provided value with the specified thread
attachOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> a -> m (Maybe a)
attachOnThread :: ThreadStorageMap a -> ThreadId -> a -> m (Maybe a)
attachOnThread ThreadStorageMap a
tsm ThreadId
tid a
ctxt = IO (Maybe a) -> m (Maybe 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 :: Int
threadAsInt = ThreadId -> Int
getThreadId ThreadId
tid
  Maybe a
old <- ThreadStorageMap a
-> (IntMap a -> (IntMap a, Maybe a)) -> IO (Maybe a)
forall a b.
ThreadStorageMap a -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe ThreadStorageMap a
tsm ((IntMap a -> (IntMap a, Maybe a)) -> IO (Maybe a))
-> (IntMap a -> (IntMap a, Maybe a)) -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ \IntMap a
m ->
    let (Maybe a
old, IntMap a
updated) = (Int -> a -> a -> a) -> Int -> a -> IntMap a -> (Maybe a, IntMap a)
forall a.
(Int -> a -> a -> a) -> Int -> a -> IntMap a -> (Maybe a, IntMap a)
I.insertLookupWithKey (\Int
_ a
n a
_ -> a
n) Int
threadAsInt a
ctxt IntMap a
m
    in (IntMap a
updated, Maybe a
old)

  case Maybe a
old of
    Maybe a
Nothing -> ThreadId -> IO () -> IO ()
addThreadFinalizer ThreadId
tid (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ThreadStorageMap a -> Int -> IO ()
forall a. ThreadStorageMap a -> Int -> IO ()
cleanUp ThreadStorageMap a
tsm Int
threadAsInt
    Just a
_ -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

  Maybe a -> IO (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
old
{-# INLINABLE attachOnThread #-}

-- | Disassociate the associated value from the current thread, returning it if it exists.
detach :: MonadIO m => ThreadStorageMap a -> m (Maybe a)
detach :: ThreadStorageMap a -> m (Maybe a)
detach ThreadStorageMap a
tsm = IO (Maybe a) -> m (Maybe 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
{-# INLINABLE detach #-}

-- | Disassociate the associated value from the specified thread, returning it if it exists.
detachFromThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> m (Maybe a)
detachFromThread :: ThreadStorageMap a -> ThreadId -> m (Maybe a)
detachFromThread ThreadStorageMap a
tsm ThreadId
tid = IO (Maybe a) -> m (Maybe 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 :: Int
threadAsInt = ThreadId -> Int
getThreadId ThreadId
tid
  ThreadStorageMap a
-> (IntMap a -> (IntMap a, Maybe a)) -> IO (Maybe a)
forall a b.
ThreadStorageMap a -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe ThreadStorageMap a
tsm ((IntMap a -> (IntMap a, Maybe a)) -> IO (Maybe a))
-> (IntMap a -> (IntMap a, Maybe a)) -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ \IntMap a
m -> (Int -> IntMap a -> IntMap a
forall a. Int -> IntMap a -> IntMap a
I.delete Int
threadAsInt IntMap a
m, Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
I.lookup Int
threadAsInt IntMap a
m)
{-# INLINABLE detachFromThread #-}

-- | Update the associated value for the current thread if it is attached.
adjust :: MonadIO m => ThreadStorageMap a -> (a -> a) -> m ()
adjust :: ThreadStorageMap a -> (a -> a) -> m ()
adjust ThreadStorageMap a
tsm a -> a
f = IO () -> m ()
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
{-# INLINABLE adjust #-}

-- | Update the associated value for the specified thread if it is attached.
adjustOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> (a -> a) -> m ()
adjustOnThread :: ThreadStorageMap a -> ThreadId -> (a -> a) -> m ()
adjustOnThread ThreadStorageMap a
tsm ThreadId
tid a -> a
f = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
  let threadAsInt :: Int
threadAsInt = ThreadId -> Int
getThreadId ThreadId
tid 
  ThreadStorageMap a -> (IntMap a -> (IntMap a, ())) -> IO ()
forall a b.
ThreadStorageMap a -> (IntMap a -> (IntMap a, b)) -> IO b
atomicModifyStripe ThreadStorageMap a
tsm ((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 Int
threadAsInt IntMap a
m, ())
{-# INLINABLE adjustOnThread #-}

-- Remove this context for thread from the map on finalization
cleanUp :: ThreadStorageMap a -> Int -> IO ()
cleanUp :: ThreadStorageMap a -> Int -> IO ()
cleanUp (ThreadStorageMap IORef (IntMap a)
tsm) Int
tid = IORef (IntMap a) -> (IntMap a -> (IntMap a, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (IntMap a)
tsm ((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 Int
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 :: ThreadStorageMap a -> IO [(Int, a)]
storedItems (ThreadStorageMap IORef (IntMap a)
tsm) = IntMap a -> [(Int, a)]
forall a. IntMap a -> [(Int, a)]
I.toList (IntMap a -> [(Int, a)]) -> IO (IntMap a) -> IO [(Int, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (IntMap a) -> IO (IntMap a)
forall a. IORef a -> IO a
readIORef IORef (IntMap a)
tsm