@@ -1,61 +1,34 @@
1−{-# LANGUAGE TupleSections #-}
21module Control.Concurrent.ThreadLocal.Internal where
32
4−import Control.Concurrent (ThreadId, myThreadId, mkWeakThreadId)
3+import Control.Concurrent (ThreadId, myThreadId)
4+import Control.Concurrent.Async
55import Control.Concurrent.MVar
6−import Control.Monad ((<=<), guard)
7−import Data.Foldable (find, toList)
8−import Data.Maybe (isJust)
9−import Data.Vector (Vector)
10−import qualified Data.Vector as Vector
11−import Data.Witherable
12−import System.Mem.Weak
6+import Data.Map (Map)
7+import qualified Data.Map.Strict as Map
138
14−data LocalValue a = LocalValue {-# UNPACK #-} !(Weak ThreadId) -- ^ ThreadId reference used as a key
15− String -- ^ String representation of ThreadId, for debugging
16− a
9+type ThreadMap a = Map ThreadId a
1710
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))
29− }
11+data ThreadLocal a = ThreadLocal { mvar :: MVar (ThreadMap a)
12+ , gcThread :: Maybe (Async ())
13+ }
3014
3115-- | Create a new thread local storage of type 'a'
3216newThreadLocal :: IO (ThreadLocal a)
33−newThreadLocal = do var <- newMVar Vector.empty
34− return (ThreadLocal var)
17+newThreadLocal = do var <- newMVar Map.empty
18+ return (ThreadLocal var Nothing)
3519
3620-- | Insert a value for the current thread
3721insertThreadLocal :: a -> ThreadLocal a -> IO ()
38−insertThreadLocal val (ThreadLocal var) = do tid <- myThreadId
39− loc <- makeLocal 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))
22+insertThreadLocal val (ThreadLocal var _) = do tid <- myThreadId
23+ modifyMVar_ var (pure . Map.insert tid val)
4524
4625-- | Fetch a value for the current thread. If none was inserted previously, return Nothing.
4726fetchThreadLocal :: ThreadLocal a -> IO (Maybe a)
48−fetchThreadLocal (ThreadLocal var) = do tid <- myThreadId
49− v <- readMVar var
50− f tid <$> wither derefLocal v
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))
27+fetchThreadLocal (ThreadLocal var _) = do tid <- myThreadId
28+ m <- readMVar var
29+ return (Map.lookup tid m)
5630
57−dumpThreadLocal :: ThreadLocal a -> IO [(Maybe ThreadId, String, a)]
58−dumpThreadLocal (ThreadLocal var) = traverse (\l@(LocalValue _ s v) -> (,s,v) <$> localThread l) . toList =<< readMVar var
31+deleteThreadLocal :: ThreadLocal a -> IO ()
32+deleteThreadLocal (ThreadLocal var _) = do tid <- myThreadId
33+ modifyMVar_ var (pure . Map.delete tid)
5934
60−printThreadLocal :: Show a => ThreadLocal a -> IO ()
61−printThreadLocal = print <=< dumpThreadLocal