Changed the way I handled the Bot monad for more concurrency.

3 files changed, 56 insertions, 44 deletions

diff --git a/Hsbot/Core.hs b/Hsbot/Core.hs
index 529e6cb..f1659e9 100644
--- a/Hsbot/Core.hs
+++ b/Hsbot/Core.hs
@@ -7,7 +7,6 @@ module Hsbot.Core
 import Control.Concurrent
 import Control.Exception (IOException, catch)
 import Control.Monad.Reader
-import Control.Monad.State
 import qualified Data.ByteString.Char8 as C
 import qualified Data.ByteString.Lazy as L
 import qualified Data.Map as M
@@ -26,6 +25,7 @@ import Hsbot.Utils
 initHsbot :: Config -> IO BotEnv
 initHsbot config = do
     chan <- newChan :: IO (Chan Message)
+    botState <- newEmptyMVar
     threadIdsMv <- newMVar []
     quitMv <- newEmptyMVar
     let hostname = configAddress config
@@ -43,7 +43,8 @@ initHsbot config = do
             handshake sCtx
             return (Just tlsenv, Just sCtx))
         else return (Nothing, Nothing)
-    return BotEnv { envHandle      = connhdl
+    return BotEnv { envBotState    = botState
+                  , envHandle      = connhdl
                   , envChan        = chan
                   , envQuitMv      = quitMv
                   , envThreadIdsMv = threadIdsMv
@@ -53,16 +54,15 @@ initHsbot config = do
 
 runHsbot :: Env IO BotStatus
 runHsbot = do
-    let bot = BotState { botPlugins  = M.empty
-                       , botHooks    = []
-                       , botChannels = []
-                       , botNickname = [] }
-    evalStateT trueRunHsbot bot
+    botNotInitialized <- asks envBotState >>= liftIO . isEmptyMVar
+    when botNotInitialized runFirstSteps
+    trueRunHsbot
   where
-    trueRunHsbot :: Bot (Env IO) BotStatus
-    trueRunHsbot = do
+    -- | Initialize the dialog with the IRC server
+    runFirstSteps :: Env IO ()
+    runFirstSteps = do
+        env <- ask
         -- First we say hello
-        env <- lift ask
         hostname <- liftIO getHostName
         let connhdl  = envHandle env
             tlsCtx   = envTLSCtx env
@@ -71,26 +71,34 @@ runHsbot = do
             channels = configChannels config
         liftIO . sendStr connhdl tlsCtx . IRC.encode $ IRC.nick nickname
         liftIO . sendStr connhdl tlsCtx . IRC.encode $ IRC.user nickname hostname "*" (configRealname config)
+        -- Then we join channels
         mapM_ (liftIO . sendStr connhdl tlsCtx . IRC.encode . IRC.joinChan) channels
+        -- Finally we set the new bot state
+        asks envBotState >>= liftIO . (flip putMVar BotState { botPlugins  = M.empty
+                                                             , botHooks    = []
+                                                             , botChannels = channels
+                                                             , botNickname = nickname })
+    -- | Run the bot itself
+    trueRunHsbot :: Env IO BotStatus
+    trueRunHsbot = do
+        env <- ask
         -- Next we spawn the reader thread
         liftIO $ debugM "Hsbot.Core" "Spawning reader thread"
+        let connhdl  = envHandle env
+            tlsCtx   = envTLSCtx env
         myOwnThreadId  <- liftIO myThreadId
-        chan <- lift $ asks envChan
-        (liftIO . forkIO $ botReader connhdl tlsCtx chan myOwnThreadId) >>= lift . addThreadIdToQuitMVar
+        chan <- asks envChan
+        (liftIO . forkIO $ botReader connhdl tlsCtx chan myOwnThreadId) >>= addThreadIdToQuitMVar
         -- Then we spawn all plugins
-        lift (asks envConfig) >>= mapM_ loadPlugin . configPlugins
+        asks envConfig >>= mapM_ loadPlugin . configPlugins
         -- Finally we spawn the main bot loop
-        bot <- get
-        finalStateMVar <- liftIO newEmptyMVar
-        (liftIO . forkIO $ runReaderT (execStateT botLoop bot >>= storeFinalState finalStateMVar) env) >>= lift . addThreadIdToQuitMVar
+        (liftIO . forkIO $ runReaderT botLoop env) >>= addThreadIdToQuitMVar
         -- We wait for the quit signal
         code <- asks envQuitMv >>= liftIO . takeMVar
         -- and we clean things up
         asks envThreadIdsMv >>= liftIO . readMVar >>= liftIO . mapM_ killThread
         -- TODO : kill plugin threads
         return code
-    storeFinalState :: MVar BotState -> BotState -> Env IO ()
-    storeFinalState finalStateMVar finalState = liftIO $ putMVar finalStateMVar finalState
 
 botReader :: Handle -> Maybe TLSCtx -> Chan Message -> ThreadId -> IO ()
 botReader _ (Just ctx) chan _ = forever $
@@ -113,16 +121,16 @@ handleIncomingStr chan str =
             writeChan chan $ IncomingMsg msg
         Nothing -> debugM "Ircd.Reader" $ "Error: couldn't decode : " ++ str     -- TODO: spam control
 
-botLoop :: Bot (Env IO) ()
+botLoop :: Env IO ()
 botLoop = forever $ do
-    chan <- lift $ asks envChan
-    hooks <- gets botHooks
+    chan <- asks envChan
     msg  <- liftIO $ readChan chan
+    hooks <- asks envBotState >>= liftIO . flip withMVar (return . botHooks)
     mapM_ (liftIO . flip writeChan msg) hooks
     case msg of
         IncomingMsg _ -> return () -- TODO parse for core commands
         OutgoingMsg outMsg -> do
-            env <- lift ask
+            env <- ask
             let connhdl  = envHandle env
                 tlsCtx   = envTLSCtx env
             liftIO $ debugM "Ircd.Reader" $ "--> " ++ show outMsg
diff --git a/Hsbot/Plugin.hs b/Hsbot/Plugin.hs
index 3191a15..6f14413 100644
--- a/Hsbot/Plugin.hs
+++ b/Hsbot/Plugin.hs
@@ -10,29 +10,32 @@ import System.Log.Logger
 
 import Hsbot.Types
 
-loadPlugin :: PluginId -> Bot (Env IO) ()
+loadPlugin :: PluginId -> Env IO ()
 loadPlugin pId = do
-    bot <- get
-    chan <- liftIO (newChan :: IO (Chan Message))
-    master <- lift $ asks envChan
-    let name = pluginName pId
-        loop = pluginEp pId
-        oldPlugins = botPlugins bot
-        pState = PluginState { pluginId     = pId
-                             , pluginChan   = chan
-                             , pluginMaster = master }
-    -- We check for unicity
-    case M.lookup name oldPlugins of
-        Just _  -> liftIO . warningM "Hsbot.Core.LoadPlugin" $ "Not loading already loaded plugin : " ++ name
-        Nothing -> do
-            liftIO . infoM "Hsbot.Core.LoadPlugin" $ "Loading plugin : " ++ name
-            env <- lift ask
-            finalStateMVar <- liftIO newEmptyMVar
-            threadId <- liftIO . forkIO $ runReaderT (execStateT loop pState >>= storeFinalState finalStateMVar) env
-            let newPlugins = M.insert name (pState, finalStateMVar, threadId) oldPlugins
-            put $ bot { botPlugins = newPlugins
-                      , botHooks   = chan : botHooks bot }
+    botMVar <- asks envBotState
+    (liftIO $ takeMVar botMVar) >>= execStateT effectivelyLoadPlugin >>= liftIO . putMVar botMVar
   where
+    effectivelyLoadPlugin :: Bot (Env IO) ()
+    effectivelyLoadPlugin = do
+        bot <- get
+        chan <- liftIO (newChan :: IO (Chan Message))
+        master <- lift $ asks envChan
+        let name = pluginName pId
+            loop = pluginEp pId
+            oldPlugins = botPlugins bot
+            pState = PluginState { pluginId     = pId
+                                 , pluginChan   = chan
+                                 , pluginMaster = master }
+        case M.lookup name oldPlugins of
+            Just _  -> liftIO . warningM "Hsbot.Core.LoadPlugin" $ "Not loading already loaded plugin : " ++ name
+            Nothing -> do
+                liftIO . infoM "Hsbot.Core.LoadPlugin" $ "Loading plugin : " ++ name
+                env <- lift ask
+                finalStateMVar <- liftIO newEmptyMVar
+                threadId <- liftIO . forkIO $ runReaderT (execStateT loop pState >>= storeFinalState finalStateMVar) env
+                let newPlugins = M.insert name (pState, finalStateMVar, threadId) oldPlugins
+                put $ bot { botPlugins = newPlugins
+                          , botHooks   = chan : botHooks bot }
     storeFinalState :: MVar PluginState -> PluginState -> Env IO ()
     storeFinalState finalStateMVar finalState = liftIO $ putMVar finalStateMVar finalState
 
diff --git a/Hsbot/Types.hs b/Hsbot/Types.hs
index b667286..14e89a1 100644
--- a/Hsbot/Types.hs
+++ b/Hsbot/Types.hs
@@ -25,7 +25,8 @@ import System.IO
 type Env = ReaderT BotEnv
 
 data BotEnv = BotEnv
-    { envHandle      :: Handle
+    { envBotState    :: MVar BotState
+    , envHandle      :: Handle
     , envChan        :: Chan Message
     , envQuitMv      :: MVar BotStatus
     , envThreadIdsMv :: MVar [ThreadId]