blob: 402c1f8ce6bf8fc193bcac9e42df296d253944ba (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
|
{-# LANGUAGE FlexibleContexts #-}
module Thermoprint.Server.Fork
( ThreadManager
, fork
, cleanup
, threadManager
) where
import Control.Monad.Reader.Class
import Control.Monad.Trans.Class
import Control.Monad.Trans.Reader (ReaderT, runReaderT)
import Control.Monad.Catch
import Control.Monad.IO.Class
import Control.Monad
import Control.Applicative
import Data.Maybe
import Data.Foldable
import Data.Map (Map)
import qualified Data.Map as Map
import Control.Concurrent
import Control.Concurrent.STM
import Control.Concurrent.STM.TVar (TVar)
import qualified Control.Concurrent.STM.TVar as T
import Control.Concurrent.STM.TSem (TSem)
import qualified Control.Concurrent.STM.TSem as S
data ThreadManager m = ThreadManager
{ fork :: m () -> m ThreadId
, cleanup :: m ()
}
threadManager :: (MonadIO m, MonadMask m) => (m () -> m ThreadId) -> m (ThreadManager m)
threadManager f = do
tVar <- newTVar Map.empty
return ThreadManager
{ fork = \act -> do
let
unregisterSelf :: MonadIO m => m ()
unregisterSelf = do
tMap <- readTVar tVar
tId <- liftIO $ myThreadId
modifyTVar' tVar $ Map.delete tId
maybeM signalTSem $ Map.lookup tId tMap
mask $ \unmask -> do
tId <- f (unmask act `finally` unregisterSelf)
modifyTVar' tVar =<< (Map.insert tId <$> newTSem 0)
return tId
, cleanup = liftIO $
mapM_ (\(tId, s) -> killThread tId >> waitTSem s) . Map.toList =<< readTVar tVar
}
where
atomically' :: MonadIO m => STM a -> m a
atomically' = liftIO . atomically
newTSem :: MonadIO m => Int -> m TSem
newTSem = atomically' . S.newTSem
waitTSem :: MonadIO m => TSem -> m ()
waitTSem = atomically' . S.waitTSem
signalTSem :: MonadIO m => TSem -> m ()
signalTSem = atomically' . S.signalTSem
newTVar :: MonadIO m => a -> m (TVar a)
newTVar = atomically' . T.newTVar
readTVar :: MonadIO m => TVar a -> m a
readTVar = atomically' . T.readTVar
modifyTVar' :: MonadIO m => TVar a -> (a -> a) -> m ()
modifyTVar' t = atomically' . T.modifyTVar t
maybeM :: Applicative m => (a -> m ()) -> Maybe a -> m ()
maybeM = maybe $ pure ()
|