Storage.hs

{-# 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 tid#) = fromIntegral (c_getThreadId 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 = liftIO (ThreadStorageMap <$> newIORef mempty)

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

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

-- | Retrieve a value if it exists for the current thread
lookup :: MonadIO m => ThreadStorageMap a -> m (Maybe a)
lookup tsm = liftIO $ do
  tid <- myThreadId
  lookupOnThread tsm tid
{-# INLINABLE lookup #-}

-- | Retrieve a value if it exists for the specified thread
lookupOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> m (Maybe a)
lookupOnThread tsm tid = liftIO $ do
  let threadAsInt = getThreadId tid
  m <- readStripe tsm
  pure $ I.lookup threadAsInt m
{-# INLINABLE lookupOnThread #-}

-- | Associate the provided value with the current thread
attach :: MonadIO m => ThreadStorageMap a -> a -> m (Maybe a)
attach tsm x = liftIO $ do
  tid <- myThreadId
  attachOnThread tsm tid x
{-# INLINABLE attach #-}

-- | Associate the provided value with the specified thread
attachOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> a -> m (Maybe a)
attachOnThread tsm tid ctxt = liftIO $ do
  let threadAsInt = getThreadId tid
  old <- atomicModifyStripe tsm $ \m ->
    let (old, updated) = I.insertLookupWithKey (\_ n _ -> n) threadAsInt ctxt m
    in (updated, old)

  case old of
    Nothing -> addThreadFinalizer tid $ cleanUp tsm threadAsInt
    Just _ -> pure ()

  pure 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 tsm = liftIO $ do
  tid <- myThreadId
  detachFromThread tsm 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 tsm tid = liftIO $ do
  let threadAsInt = getThreadId tid
  atomicModifyStripe tsm $ \m -> (I.delete threadAsInt m, I.lookup threadAsInt m)
{-# INLINABLE detachFromThread #-}

-- | Update the associated value for the current thread if it is attached.
adjust :: MonadIO m => ThreadStorageMap a -> (a -> a) -> m ()
adjust tsm f = liftIO $ do
  tid <- myThreadId
  adjustOnThread tsm tid 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 tsm tid f = liftIO $ do
  let threadAsInt = getThreadId tid 
  atomicModifyStripe tsm $ \m -> (I.adjust f threadAsInt m, ())
{-# INLINABLE adjustOnThread #-}

-- Remove this context for thread from the map on finalization
cleanUp :: ThreadStorageMap a -> Int -> IO ()
cleanUp (ThreadStorageMap tsm) tid = atomicModifyIORef tsm $ \m -> 
  (I.delete tid 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 tsm) = I.toList <$> readIORef tsm