@@ -0,0 +1,58 @@
1+module Control.Concurrent.ThreadLocal.Internal where
2+
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
8+import System.Mem.Weak
9+
10+type ThreadName = StableName ThreadId
11+
12+data LocalVal a = LocalVal !(Weak ThreadId) !a
13+
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'
19+
20+newtype ThreadLocal a = ThreadLocal { values :: IORef (Map ThreadName (LocalVal a))
21+ }
22+
23+-- | Create a new thread local storage of type 'a'
24+newThreadLocal :: IO (ThreadLocal a)
25+newThreadLocal = do ref <- newIORef Map.empty
26+ return (ThreadLocal ref)
27+
28+-- | Insert a value for the current thread
29+insertThreadLocal :: 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',())
36+
37+-- | Fetch a value for the current thread. If none was inserted previously, return Nothing.
38+fetchThreadLocal :: 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, ())