concurrency & dyre fixes for server spec

1 files changed, 22 insertions, 11 deletions

diff --git a/server/src/Thermoprint/Server.hs b/server/src/Thermoprint/Server.hs
index 678d056..4559414 100644
--- a/server/src/Thermoprint/Server.hs
+++ b/server/src/Thermoprint/Server.hs
@@ -23,6 +23,9 @@ import qualified Config.Dyre as Dyre
 import Data.Map (Map)
 import qualified Data.Map as Map
 
+import Data.Set (Set)
+import qualified Data.Set as Set
+
 import Data.Maybe (maybe)
 import Data.Foldable (mapM_, forM_, foldlM)
 
@@ -34,7 +37,7 @@ import Control.Monad.Reader
 import Control.Monad.IO.Class
 import Control.Monad.Morph
 import Control.Category
-import Control.Monad.Catch (MonadMask)
+import Control.Monad.Catch (MonadMask(mask), finally)
 import Prelude hiding (id, (.))
 
 import qualified Control.Monad as M
@@ -56,12 +59,16 @@ import Database.Persist.Sql (runMigrationSilent, ConnectionPool, runSqlPool)
 
 import Thermoprint.API (thermoprintAPI, PrinterId)
 
+import Thermoprint.Server.Fork
+
 import Thermoprint.Server.Database
 import Thermoprint.Server.Printer
 import Thermoprint.Server.Queue
 import qualified Thermoprint.Server.API as API (thermoprintServer)
 import Thermoprint.Server.API hiding (thermoprintServer)
 
+import Debug.Trace
+
 -- | Compile-time configuration for 'thermoprintServer'
 data Config m = Config { dyreError     :: Maybe String -- ^ Set by 'Dyre' -- sent to log as an error
                        , warpSettings  :: Warp.Settings -- ^ Configure 'Warp's behaviour
@@ -107,21 +114,25 @@ thermoprintServer :: ( MonadLoggerIO m
                      , MonadReader ConnectionPool m
                      , MonadResourceBase m
                      , MonadMask m
-                     ) => (m :~> IO) -- ^ 'dyre' controls the base of the monad-transformer-stack ('IO') but we let the user specify much of the rest of it (we handle 'ResourceT' ourselves, since we need it to fork properly). Therefore we require a specification of how to collapse the stack.
+                     ) => Bool -- ^ Invoke 'dyre' to look for and attempt to compile custom configurations (pass 'False' iff testing)
+                     -> (m :~> IO) -- ^ 'dyre' controls the base of the monad-transformer-stack ('IO') but we let the user specify much of the rest of it (we handle 'ResourceT' ourselves, since we need it to fork properly). Therefore we require a specification of how to collapse the stack.
                      -> ResourceT m (Config (ResourceT m)) -> IO ()
 -- ^ Run the server
-thermoprintServer io = Dyre.wrapMain $ Dyre.defaultParams
+thermoprintServer dyre io = Dyre.wrapMain $ Dyre.defaultParams
   { Dyre.projectName = "thermoprint-server"
   , Dyre.realMain    = realMain
   , Dyre.showError   = flip (\msg -> fmap (\cfg -> cfg { dyreError = Just msg }))
+  , Dyre.configCheck = dyre
   }
     where
       realMain cfg = unNat (io . Nat runResourceT) $ do
-        Config{..} <- cfg
-        maybe (return ()) ($(logErrorS) "Dyre" . T.pack) dyreError
-        mapM_ ($(logWarnS) "DB") =<< runSqlPool (runMigrationSilent migrateAll) =<< ask
-        forM_ printers $ resourceForkIO . runPrinter
-        let
-          runQM' (queueManagers -> QMConfig qm nat) printer = unNat nat $ runQM qm printer
-        Map.foldrWithKey (\k p a -> resourceForkIO (runQM' k p) >> a) (return ()) printers
-        liftIO . Warp.runSettings warpSettings . serve thermoprintAPI . flip enter API.thermoprintServer =<< handlerNat printers
+        tMgr <- threadManager resourceForkIO
+        flip finally (cleanup tMgr) $ do
+          Config{..} <- cfg
+          maybe (return ()) ($(logErrorS) "Dyre" . T.pack) dyreError
+          mapM_ ($(logWarnS) "DB") =<< runSqlPool (runMigrationSilent migrateAll) =<< ask
+          forM_ printers $ fork tMgr . runPrinter
+          let
+            runQM' (queueManagers -> QMConfig qm nat) printer = unNat nat $ runQM qm printer
+          mapM_ (fork tMgr . uncurry runQM') $ Map.toList printers
+          liftIO . Warp.runSettings warpSettings . serve thermoprintAPI . flip enter API.thermoprintServer =<< handlerNat printers