thread-utils-context 0.2.0.0 → 0.3.0.0
raw patch · 5 files changed
+160/−112 lines, 5 filesdep +hspecdep +hspec-expectationsdep −criteriondep −mtl
Dependencies added: hspec, hspec-expectations
Dependencies removed: criterion, mtl
Files
- ChangeLog.md +0/−5
- bench/Bench.hs +0/−19
- src/Control/Concurrent/Thread/Storage.hs +113/−52
- test/Spec.hs +36/−16
- thread-utils-context.cabal +11/−20
ChangeLog.md view
@@ -1,8 +1,3 @@ # Changelog for thread-utils-context -## 0.2.0.0--- Significant performance improvements-- 'storedItems' now returns both keys and values.- ## Unreleased changes
− bench/Bench.hs
@@ -1,19 +0,0 @@-import Criterion-import Criterion.Main-import Control.Concurrent-import Control.Concurrent.Thread.Storage -import Control.Monad.Reader--main :: IO ()-main = do- tsm <- newThreadStorageMap- defaultMain- [ bench "attach" $ whnfIO $ do- attach tsm ()- -- , bench "myThreadId" $ whnfIO myThreadId- -- , bench "hashed getThreadId" $ whnfIO $ do- -- tid <- myThreadId- -- pure $ threadHash $ getThreadId tid- -- , bench "local" $ whnfIO $ runReaderT (local (const ()) pure ()) ()- ]- -- print =<< storedItemsCount tsm
src/Control/Concurrent/Thread/Storage.hs view
@@ -1,6 +1,8 @@ {-# LANGUAGE MagicHash #-} {-# LANGUAGE UnliftedFFITypes #-} {-# LANGUAGE UnboxedTuples #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE Strict #-} -- | 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@@ -18,13 +20,10 @@ -- 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.+-- 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'@@ -33,6 +32,9 @@ -- * Retrieve values from a 'ThreadStorageMap' , lookup , lookupOnThread+ -- * Update values in a 'ThreadStorageMap'+ , update+ , updateOnThread -- * Associate values with a thread in a 'ThreadStorageMap' , attach , attachOnThread@@ -44,114 +46,158 @@ , 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 ( void )+import Control.Monad ( when, void ) import Control.Monad.IO.Class-import Data.IORef-import GHC.IO (IO(..))+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 :: ThreadId# -> CInt+foreign import ccall unsafe "rts_getThreadId" c_getThreadId :: Addr# -> CULLong +numStripes :: Int+numStripes = 32+ getThreadId :: ThreadId -> Int-getThreadId (ThreadId tid#) = fromIntegral (c_getThreadId tid#)-{-# INLINE getThreadId #-}+getThreadId (ThreadId tid#) = fromIntegral (c_getThreadId (unsafeCoerce# tid#))++stripeHash :: Int -> Int+stripeHash = (`mod` numStripes) . abs++readStripe :: ThreadStorageMap a -> ThreadId -> IO (I.IntMap a)+readStripe (ThreadStorageMap arr#) t = IO $ \s -> readArray# arr# tid# s+ where+ (I# tid#) = stripeHash $ getThreadId t++atomicModifyStripe :: ThreadStorageMap a -> Int -> (I.IntMap a -> (I.IntMap a, b)) -> IO b+atomicModifyStripe (ThreadStorageMap arr#) tid f = IO $ \s -> go s+ where+ (I# stripe#) = stripeHash tid+ go s = case readArray# arr# stripe# s of+ (# s1, intMap #) ->+ let (updatedIntMap, result) = f intMap + in case casArray# arr# stripe# intMap updatedIntMap s1 of+ (# s2, outcome, old #) -> case outcome of+ 0# -> (# s2, result #)+ 1# -> go s2+ _ -> error "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.-newtype ThreadStorageMap a = ThreadStorageMap (IORef (I.IntMap a))+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 = 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 #-}+newThreadStorageMap = liftIO $ IO $ \s -> case newArray# numStripes# mempty s of+ (# s1, ma #) -> (# s1, ThreadStorageMap ma #)+ where+ (I# numStripes#) = numStripes -- | 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+ m <- readStripe tsm tid pure $ I.lookup threadAsInt m-{-# INLINABLE lookupOnThread #-}+ where + threadAsInt = getThreadId tid --- | Associate the provided value with the current thread+-- | 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 tsm x = liftIO $ do tid <- myThreadId attachOnThread tsm tid x-{-# INLINABLE attach #-} --- | Associate the provided value with the specified thread+-- | 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 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 #-}+attachOnThread tsm tid ctxt = + updateOnThread tsm tid (\prev -> (Just ctxt, prev)) -- | 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 #-}+ updateOnThread tsm tid (\prev -> (Nothing, 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 tsm tid f = liftIO $ do+ -- ^ We mask here in order to ensure that the finalizer will always be created+ (isNewThreadEntry, result) <- atomicModifyStripe tsm threadAsInt $ \m -> + let (resultWithNewThreadDetection, m') = + I.alterF + (\x -> case f x of+ (x', y) -> ((isNothing x && isJust x', y), x')+ ) + threadAsInt+ m+ in (m', resultWithNewThreadDetection)+ when isNewThreadEntry $ do+ addThreadFinalizer tid $ cleanUp tsm threadAsInt+ pure result+ where + threadAsInt = getThreadId tid++update :: MonadIO m => ThreadStorageMap a -> (Maybe a -> (Maybe a, b)) -> m b+update tsm f = liftIO $ do+ tid <- myThreadId+ updateOnThread tsm tid f+ -- | 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 #-}+ atomicModifyStripe tsm threadAsInt $ \m -> (I.adjust f threadAsInt m, ())+ where + threadAsInt = getThreadId tid -- 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, ())+cleanUp tsm tid = do+ atomicModifyStripe tsm tid $ \m -> + (I.delete tid m, ()) -- | List thread ids with live entries in the 'ThreadStorageMap'. -- @@ -159,4 +205,19 @@ -- 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+storedItems tsm = do+ stripes <- mapM (stripeByIndex tsm) [0..(numStripes - 1)]+ pure $ concatMap I.toList stripes+ where+ stripeByIndex :: ThreadStorageMap a -> Int -> IO (I.IntMap a)+ stripeByIndex (ThreadStorageMap arr#) (I# i#) = IO $ \s -> readArray# arr# i# 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 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
test/Spec.hs view
@@ -4,22 +4,42 @@ import Control.Concurrent.MVar import Control.Concurrent.Thread.Storage import Control.Monad+import Data.List hiding (lookup)+import Test.Hspec+import Prelude hiding (lookup) main :: IO ()-main = do- mv <- newEmptyMVar- tsm <- newThreadStorageMap- replicateM_ 20 $ do- forkIO $ do- myThreadId >>= print- attach tsm ()- readMVar mv- threadDelay 2_000_000- print =<< storedItems tsm- putMVar mv ()- threadDelay 2_000_000- performGC- print =<< storedItems tsm-- +main = hspec $ do + describe "cleanup" $ do+ it "works" $ do+ mv <- newEmptyMVar+ tsm <- newThreadStorageMap+ replicateM_ 100000 $ do+ forkIO $ do+ attach tsm ()+ readMVar mv+ threadDelay 10_000_000+ putMVar mv ()+ threadDelay 10_000_000+ performGC+ threadDelay 10_000_000+ thingsStillInStorage <- storedItems tsm+ sort thingsStillInStorage `shouldBe` []+ it "doesn't happen while a thread is still alive" $ do+ tsm <- newThreadStorageMap+ mv <- newEmptyMVar+ resultVar <- newEmptyMVar+ forkIO $ do+ attach tsm ()+ readMVar mv+ putMVar resultVar =<< lookup tsm+ threadDelay 5_000_000+ performGC+ putMVar mv ()+ result <- readMVar resultVar+ result `shouldBe` Just ()+ performGC+ threadDelay 5_000_000+ items <- storedItems tsm+ items `shouldBe` []
thread-utils-context.cabal view
@@ -1,11 +1,11 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.34.4.+-- This file has been generated from package.yaml by hpack version 0.35.2. -- -- see: https://github.com/sol/hpack name: thread-utils-context-version: 0.2.0.0+version: 0.3.0.0 synopsis: Garbage-collected thread local storage description: Please see the README on GitHub at <https://github.com/iand675/thread-utils-context#readme> category: Concurrency@@ -25,6 +25,11 @@ type: git location: https://github.com/iand675/thread-utils +flag debug+ description: Whether to enable some additional hooks to debug issues+ manual: True+ default: False+ library exposed-modules: Control.Concurrent.Thread.Storage@@ -38,6 +43,8 @@ , ghc-prim , thread-utils-finalizers default-language: Haskell2010+ if flag(debug)+ cpp-options: -DDEBUG_HOOKS test-suite thread-utils-context-test type: exitcode-stdio-1.0@@ -51,24 +58,8 @@ base >=4.7 && <5 , containers , ghc-prim- , thread-utils-context- , thread-utils-finalizers- default-language: Haskell2010--benchmark thread-utils-context-benchmarks- type: exitcode-stdio-1.0- main-is: Bench.hs- other-modules:- Paths_thread_utils_context- hs-source-dirs:- bench- ghc-options: -O2 -rtsopts -threaded -with-rtsopts=-N- build-depends:- base >=4.7 && <5- , containers- , criterion- , ghc-prim- , mtl+ , hspec+ , hspec-expectations , thread-utils-context , thread-utils-finalizers default-language: Haskell2010