@@ -1,58 +1,62 @@
1+{-# LANGUAGE TupleSections #-}
12module Control.Concurrent.ThreadLocal.Internal where
23
3−import Control.Concurrent (ThreadId, myThreadId)
4−import Data.IORef
5−import Data.Map (Map)
6−import qualified Data.Map.Strict as Map
7−import System.Mem.StableName
4+import Control.Concurrent (ThreadId, myThreadId, mkWeakThreadId)
5+import Control.Concurrent.MVar
6+import Control.Monad ((<=<), guard)
7+import Data.Foldable (toList)
8+import Data.Maybe (isJust)
9+import Data.Vector (Vector)
10+import qualified Data.Vector as Vector
11+import Data.Witherable
812import System.Mem.Weak
913
10−type ThreadName = StableName ThreadId
14+data LocalValue a = LocalValue {-# UNPACK #-} !(Weak ThreadId) -- ^ ThreadId reference used as a key
15+ String -- ^ String representation of ThreadId, for debugging
16+ a
1117
12−data LocalVal a = LocalVal !(Weak ThreadId) !a
18+localThread :: LocalValue a -> IO (Maybe ThreadId)
19+localThread (LocalValue weak _ _) = deRefWeak weak
1320
14−isSameThread :: ThreadId -> Weak ThreadId -> IO Bool
15−isSameThread tid wid = do mtid' <- deRefWeak wid
16− return $ case mtid' of
17− Nothing -> False
18− Just tid' -> tid == tid'
21+derefLocal :: LocalValue a -> IO (Maybe (ThreadId, a))
22+derefLocal (LocalValue weak _ a) = fmap (,a) <$> deRefWeak weak
1923
20−newtype ThreadLocal a = ThreadLocal { values :: IORef (Map ThreadName (LocalVal a))
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))
2129 }
2230
2331-- | Create a new thread local storage of type 'a'
2432newThreadLocal :: IO (ThreadLocal a)
25−newThreadLocal = do ref <- newIORef Map.empty
26− return (ThreadLocal ref)
33+newThreadLocal = do var <- newMVar Vector.empty
34+ return (ThreadLocal var)
2735
2836-- | Insert a value for the current thread
2937insertThreadLocal :: a -> ThreadLocal a -> IO ()
30−insertThreadLocal v (ThreadLocal ref) = do tid <- myThreadId
31− name <- makeStableName tid
32− wid <- mkWeakPtr tid Nothing
33− atomicModifyIORef' ref $ \m ->
34− let m' = Map.insert name (LocalVal wid v) m
35− in (m',())
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))
3645
3746-- | Fetch a value for the current thread. If none was inserted previously, return Nothing.
3847fetchThreadLocal :: ThreadLocal a -> IO (Maybe a)
39−fetchThreadLocal (ThreadLocal ref) = do tid <- myThreadId
40− name <- makeStableName tid
41− m <- readIORef ref
42− case Map.lookup name m of
43− Just (LocalVal wid v) ->
44− do st <- isSameThread tid wid
45− if st
46− then return (Just v)
47− -- Account for the fact that a string representation of
48− -- a ThreadId can result in a collision.
49− -- If one occurs, remove the old data.
50− else Nothing <$ removeThreadLocal (ThreadLocal ref)
51− Nothing -> return Nothing
52−
53−-- | Remove the currently held value for the current thread.
54−removeThreadLocal :: ThreadLocal a -> IO ()
55−removeThreadLocal (ThreadLocal ref) = do tid <- myThreadId
56− name <- makeStableName tid
57− atomicModifyIORef' ref $ \m ->
58− (Map.delete name m, ())
48+fetchThreadLocal (ThreadLocal var) = do tid <- myThreadId
49+ v <- readMVar var
50+ f tid <$> wither derefLocal v
51+ where f tid vec = do ((tid',v),_) <- Vector.uncons vec
52+ guard (tid == tid')
53+ pure v
54+
55+gcThreadLocal :: ThreadLocal a -> IO ()
56+gcThreadLocal (ThreadLocal var) = modifyMVar_ var (filterA (fmap isJust . localThread))
57+
58+dumpThreadLocal :: ThreadLocal a -> IO [(Maybe ThreadId, String, a)]
59+dumpThreadLocal (ThreadLocal var) = traverse (\l@(LocalValue _ s v) -> (,s,v) <$> localThread l) . toList =<< readMVar var
60+
61+printThreadLocal :: Show a => ThreadLocal a -> IO ()
62+printThreadLocal = print <=< dumpThreadLocal