1 files changed, 81 insertions, 0 deletions

diff --git a/server/src/Thermoprint/Server/Fork.hs b/server/src/Thermoprint/Server/Fork.hs
new file mode 100644
index 0000000..402c1f8
--- /dev/null
+++ b/server/src/Thermoprint/Server/Fork.hs
@@ -0,0 +1,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 ()