packages feed

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 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