10965a0 · typeable/thread-local

git.iko.soy/typeable/thread-local

on master help sign in

Make better thread local storage lol

10965a0Nick3 years ago

example/Main.hs+4−4

compare with master

unifiedsplit

@@ -9,9 +9,9 @@ import Control.Monad

9import System.Random

11main :: IO ()

12−main = do tls <- newThreadLocal

12+main = do tls <- newThreadLocalWithGC 30

13− replicateConcurrently_ 200 (f tls)

13+ replicateConcurrently_ 200000 (f tls)

14− printThreadLocal tls

14+-- printThreadLocal tls

16f :: ThreadLocal Word -> IO ()

17f tls = do v :: Word <- randomIO

@@ -19,4 +19,4 @@ f tls = do v :: Word <- randomIO

19 insertThreadLocal v tls

20 !v' <- fetchThreadLocal tls

21 unless (Just v == v') $

22− putStrLn ("Ooops! " <> show tid <> " expected " <> show v <> " while the fetched value is " <> show v')

22+ putStrLn ("Ooops! " <> show tid <> " expected " <> show v <> " while the fetched value is " <> show v')

src/Control/Concurrent/ThreadLocal.hs+7−3

@@ -1,10 +1,14 @@

1module Control.Concurrent.ThreadLocal ( ThreadLocal -- exported opaque

2 , newThreadLocal

3+ , newThreadLocalWithGC

3 , insertThreadLocal

4 , insertThreadLocal

4 , fetchThreadLocal

5 , fetchThreadLocal

5− , gcThreadLocal

6+ , deleteThreadLocal

6− , dumpThreadLocal

7− , printThreadLocal

8 ) where

7 ) where

10import Control.Concurrent.ThreadLocal.Internal

9import Control.Concurrent.ThreadLocal.Internal

10+import Control.Concurrent.ThreadLocal.GC

11+

12+newThreadLocalWithGC :: Word -> IO (ThreadLocal a)

13+newThreadLocalWithGC interval = do tls <- newThreadLocal

14+ addGC tls interval

src/Control/Concurrent/ThreadLocal/GC.hsnew file+46

@@ -0,0 +1,46 @@

1+module Control.Concurrent.ThreadLocal.GC where

3+import Control.Concurrent.ThreadLocal.Internal

5+import Control.Concurrent.Async

6+import Control.Concurrent.MVar

7+import Data.Either

8+import qualified Data.Map as Map

9+import GHC.Conc

10+import System.Mem.Weak

11+

12+-- | Check whether a thread is still alive.

13+isThreadAlive :: ThreadId -> IO Bool

14+isThreadAlive tid = isAlive <$> threadStatus tid

15+ where isAlive ThreadRunning = True

16+ isAlive (ThreadBlocked _) = True

17+ isAlive _ = False

18+

19+-- | Sort alive from dead

20+splitAlive :: [ThreadId] -> IO ([ThreadId], [ThreadId])

21+splitAlive = fmap partitionEithers

22+ . traverse (\t -> s t <$> isThreadAlive t)

23+ where s x b = if b then Left x else Right x

24+

25+collectGarbage :: ThreadMap a -> IO (ThreadMap a)

26+collectGarbage m = do (_live, dead) <- splitAlive (Map.keys m)

27+ pure (foldr Map.delete m dead)

28+

29+gcLoop :: Weak (MVar (ThreadMap a)) -> Word -> IO ()

30+gcLoop wvar t = do threadDelay (fromIntegral t * 1000000)

31+ var' <- deRefWeak wvar

32+ whenJust var' $ \var ->

33+ do modifyMVar_ var collectGarbage

34+ gcLoop wvar t

35+

36+whenJust :: Monad m => Maybe a -> (a -> m ()) -> m ()

37+whenJust (Just v) f = f v

38+whenJust Nothing _ = pure ()

39+{-# INLINE whenJust #-}

40+

41+addGC :: ThreadLocal a -> Word -> IO (ThreadLocal a)

42+addGC (ThreadLocal var _) interval =

43+ do wvar <- mkWeakMVar var (pure ())

44+ a <- async (gcLoop wvar interval)

45+ labelThread (asyncThreadId a) "ThreadLocal GC Loop"

46+ pure (ThreadLocal var (Just a))

src/Control/Concurrent/ThreadLocal/Internal.hs+18−45

@@ -1,61 +1,34 @@

1−{-# LANGUAGE TupleSections #-}

2module Control.Concurrent.ThreadLocal.Internal where

1module Control.Concurrent.ThreadLocal.Internal where

4−import Control.Concurrent (ThreadId, myThreadId, mkWeakThreadId)

3+import Control.Concurrent (ThreadId, myThreadId)

4+import Control.Concurrent.Async

5import Control.Concurrent.MVar

6−import Control.Monad ((<=<), guard)

6+import Data.Map (Map)

7−import Data.Foldable (find, toList)

8−import Data.Maybe (isJust)

9−import Data.Vector (Vector)

10−import qualified Data.Vector as Vector

7+import qualified Data.Map.Strict as Map

11−import Data.Witherable

12−import System.Mem.Weak

14−data LocalValue a = LocalValue {-# UNPACK #-} !(Weak ThreadId) -- ^ ThreadId reference used as a key

9+type ThreadMap a = Map ThreadId a

15− String -- ^ String representation of ThreadId, for debugging

16− a

18−localThread :: LocalValue a -> IO (Maybe ThreadId)

19−localThread (LocalValue weak _ _) = deRefWeak weak

20−

21−derefLocal :: LocalValue a -> IO (Maybe (ThreadId, a))

22−derefLocal (LocalValue weak _ a) = fmap (,a) <$> deRefWeak weak

23−

24−makeLocal :: ThreadId -> a -> IO (LocalValue a)

25−makeLocal tid a = do weak <- mkWeakThreadId tid

26− pure (LocalValue weak (show tid) a)

27−

28−newtype ThreadLocal a = ThreadLocal { values :: MVar (Vector (LocalValue a))

11+data ThreadLocal a = ThreadLocal { mvar :: MVar (ThreadMap a)

12+ , gcThread :: Maybe (Async ())

29− }

13+ }

31-- | Create a new thread local storage of type 'a'

15-- | Create a new thread local storage of type 'a'

32newThreadLocal :: IO (ThreadLocal a)

16newThreadLocal :: IO (ThreadLocal a)

33−newThreadLocal = do var <- newMVar Vector.empty

17+newThreadLocal = do var <- newMVar Map.empty

34− return (ThreadLocal var)

18+ return (ThreadLocal var Nothing)

36-- | Insert a value for the current thread

20-- | Insert a value for the current thread

37insertThreadLocal :: a -> ThreadLocal a -> IO ()

21insertThreadLocal :: a -> ThreadLocal a -> IO ()

38−insertThreadLocal val (ThreadLocal var) = do tid <- myThreadId

22+insertThreadLocal val (ThreadLocal var _) = do tid <- myThreadId

39− loc <- makeLocal tid val

23+ modifyMVar_ var (pure . Map.insert tid val)

40− modifyMVar_ var $ \v ->

41− do ts <- filterA (\l -> do tid' <- localThread l

42− pure (isJust tid' && Just tid /= tid')

43− ) (toList v)

44− pure (Vector.fromList (loc:ts))

46-- | Fetch a value for the current thread. If none was inserted previously, return Nothing.

25-- | Fetch a value for the current thread. If none was inserted previously, return Nothing.

47fetchThreadLocal :: ThreadLocal a -> IO (Maybe a)

26fetchThreadLocal :: ThreadLocal a -> IO (Maybe a)

48−fetchThreadLocal (ThreadLocal var) = do tid <- myThreadId

27+fetchThreadLocal (ThreadLocal var _) = do tid <- myThreadId

49− v <- readMVar var

28+ m <- readMVar var

50− f tid <$> wither derefLocal v

29+ return (Map.lookup tid m)

51− where f tid vec = do (_,v) <- find ((== tid) . fst) vec

52− pure v

53−

54−gcThreadLocal :: ThreadLocal a -> IO ()

55−gcThreadLocal (ThreadLocal var) = modifyMVar_ var (filterA (fmap isJust . localThread))

57−dumpThreadLocal :: ThreadLocal a -> IO [(Maybe ThreadId, String, a)]

31+deleteThreadLocal :: ThreadLocal a -> IO ()

58−dumpThreadLocal (ThreadLocal var) = traverse (\l@(LocalValue _ s v) -> (,s,v) <$> localThread l) . toList =<< readMVar var

32+deleteThreadLocal (ThreadLocal var _) = do tid <- myThreadId

33+ modifyMVar_ var (pure . Map.delete tid)

60−printThreadLocal :: Show a => ThreadLocal a -> IO ()

61−printThreadLocal = print <=< dumpThreadLocal

src/Control/Concurrent/ThreadLocal/Unsafe.hsdeleted−9

@@ -1,9 +0,0 @@

1−module Control.Concurrent.ThreadLocal.Unsafe where

2−

3−--import Control.Concurrent.ThreadLocal.Internal

4−--import Data.Map (Map)

5−--import qualified Data.Map (Map)

6−

7−

8−--dumpThreadLocal :: ThreadLocal a -> IO (Map ThreadId a)

9−--dumpThreadLocal (ThreadLocal ref) = do m <-

thread-local.cabal+4−3

@@ -18,10 +18,11 @@ common warnings

18library

19 import: warnings

20 exposed-modules: Control.Concurrent.ThreadLocal

21− Control.Concurrent.ThreadLocal.Unsafe

22− other-modules: Control.Concurrent.ThreadLocal.Internal

21+ other-modules: Control.Concurrent.ThreadLocal.GC

22+ Control.Concurrent.ThreadLocal.Internal

23 build-depends: base >=4.16 && < 5,

24− vector >= 0.13,

24+ async >= 2.2,

25+ containers >= 0.6,

25 witherable

26 witherable

26 hs-source-dirs: src

27 hs-source-dirs: src

27 default-language: Haskell2010

28 default-language: Haskell2010