Fixed the plugin termination.

The master of a plugin will now wait for the plugin it kills to terminate.

6 files changed, 32 insertions, 20 deletions

diff --git a/Hsbot/Irc/Command.hs b/Hsbot/Irc/Command.hs
index 242e12d..51c2187 100644
--- a/Hsbot/Irc/Command.hs
+++ b/Hsbot/Irc/Command.hs
@@ -12,6 +12,7 @@ import Data.Maybe
 import Hsbot.Irc.Message
 import Hsbot.Irc.Plugin
 import Hsbot.Irc.Types
+import Hsbot.Types
 
 -- | Registers a plugin's command
 registerCommand :: String -> String -> IrcBot ()
@@ -34,19 +35,19 @@ unregisterCommand cmd pluginName' = do
     put $ ircBot { ircBotCommands = newCmds }
 
 -- | Processes an internal command
-processInternalCommand :: IrcBotMsg -> IrcBot (Bool)
+processInternalCommand :: IrcBotMsg -> IrcBot (BotStatus)
 processInternalCommand (IntIrcCmd ircCmd)
   | ircCmdTo ircCmd == "CORE"   = processCoreCommand ircCmd
   | otherwise = do
       plugins <- gets ircBotPlugins
       case M.lookup (ircCmdTo ircCmd) plugins of
-          Just (plugin, _) -> sendToPlugin (IntIrcCmd ircCmd) plugin
+          Just (plugin, _, _) -> sendToPlugin (IntIrcCmd ircCmd) plugin
           Nothing          -> return ()
-      return False
-processInternalCommand _ = return (False)
+      return BotContinue
+processInternalCommand _ = return (BotContinue)
 
 -- | Processes a core command
-processCoreCommand :: IrcCmd -> IrcBot (Bool)
+processCoreCommand :: IrcCmd -> IrcBot (BotStatus)
 processCoreCommand ircCmd = do
     let command' = ircCmdCmd ircCmd
         originalRequest = ircCmdBotMsg ircCmd
@@ -58,7 +59,9 @@ processCoreCommand ircCmd = do
         "UNREGISTER" -> unregisterCommand (ircCmdMsg ircCmd) (ircCmdFrom ircCmd)
         "UPDATE"     -> processUpdateCommand ircCmd
         _            -> return ()
-    return $ command' == "REBOOT"
+    if command' == "REBOOT"
+      then return BotReboot
+      else return BotContinue
 
 -- | Process an update command
 processUpdateCommand :: IrcCmd -> IrcBot ()
diff --git a/Hsbot/Irc/Core.hs b/Hsbot/Irc/Core.hs
index 229fef8..51032a8 100644
--- a/Hsbot/Irc/Core.hs
+++ b/Hsbot/Irc/Core.hs
@@ -110,7 +110,7 @@ ircBotLoop = forever $ do
         IntIrcCmd intIrcCmd -> do
             reboot <- processInternalCommand $ IntIrcCmd intIrcCmd
             reportUpdate
-            if reboot
+            if reboot == BotReboot
               then processRebootCommand
               else return ()
   where
@@ -131,9 +131,9 @@ dispatchMessage (InIrcMsg inIrcMsg) = do
         let key         = tail . head $ words getMsgContent
             pluginNames = fromMaybe [] $ M.lookup key cmds
             plugins'    = fromMaybe [] $ mapM (flip M.lookup plugins) pluginNames
-        in mapM_ (sendRunCommand (tail getMsgContent) . fst) plugins'
+        in mapM_ (sendRunCommand (tail getMsgContent) . first) plugins'
       else
-        mapM_ (sendToPlugin (InIrcMsg inIrcMsg) . fst) (M.elems plugins)
+        mapM_ (sendToPlugin (InIrcMsg inIrcMsg) . first) (M.elems plugins)
   where
     isPluginCommand :: IrcConfig -> Bool
     isPluginCommand config =
diff --git a/Hsbot/Irc/Plugin.hs b/Hsbot/Irc/Plugin.hs
index 28026aa..2c8e84b 100644
--- a/Hsbot/Irc/Plugin.hs
+++ b/Hsbot/Irc/Plugin.hs
@@ -52,11 +52,12 @@ loadIrcPlugin pluginName = do
     case M.lookup pluginName oldPlugins of
         Just _  -> return ()
         Nothing -> do
-            threadId <- liftIO $ forkIO (entryPoint pluginChan masterChan)
+            mvar <- liftIO newEmptyMVar
+            threadId <- liftIO . forkIO $ finally (entryPoint pluginChan masterChan) (putMVar mvar ())
             let plugin        = IrcPluginState { ircPluginName       = pluginName
                                                , ircPluginChan       = pluginChan
                                                , ircPluginMasterChan = masterChan }
-                newPlugins    = M.insert pluginName (plugin, threadId) oldPlugins
+                newPlugins    = M.insert pluginName (plugin, mvar, threadId) oldPlugins
                 newResumeData = M.insert "PLUGINS" (show $ M.keys newPlugins) oldResumeData
             put $ ircbot { ircBotPlugins    = newPlugins
                          , ircBotResumeData = newResumeData }
@@ -67,7 +68,7 @@ listPlugins originalRequest dest = do
     plugins <- gets ircBotPlugins
     let listing = unwords $ M.keys plugins
     case M.lookup dest plugins of
-        Just (plugin, _) -> sendToPlugin (IntIrcCmd $ IrcCmd "ANSWER" "CORE" dest listing originalRequest) plugin
+        Just (plugin, _, _) -> sendToPlugin (IntIrcCmd $ IrcCmd "ANSWER" "CORE" dest listing originalRequest) plugin
         Nothing          -> return ()
 
 -- | Unloads a plugin
@@ -87,9 +88,10 @@ killIrcPlugin name = do
     let oldPlugins = ircBotPlugins ircbot
     -- We check if the plugin exists
     case M.lookup name oldPlugins of
-        Just (_, threadId) -> do
+        Just (_, mvar, threadId) -> do
             let newPlugins = M.delete name oldPlugins
             liftIO $ throwTo threadId UserInterrupt
             put $ ircbot { ircBotPlugins = newPlugins }
+            liftIO $ takeMVar mvar
         Nothing            -> return ()
 
diff --git a/Hsbot/Irc/Types.hs b/Hsbot/Irc/Types.hs
index fe70ea4..4b75085 100644
--- a/Hsbot/Irc/Types.hs
+++ b/Hsbot/Irc/Types.hs
@@ -3,6 +3,7 @@ module Hsbot.Irc.Types
     , IrcBotState (..)
     , IrcServer
     , IrcServerState (..)
+    , first
     ) where
 
 import Control.Concurrent
@@ -22,7 +23,7 @@ type IrcBot = StateT IrcBotState IO
 -- | An Ircbot state
 data IrcBotState = IrcBotState
     { ircBotStartTime            :: UTCTime                 -- the bot's uptime
-    , ircBotPlugins              :: M.Map String (IrcPluginState, ThreadId) -- Loaded plugins
+    , ircBotPlugins              :: M.Map String (IrcPluginState, MVar (), ThreadId) -- Loaded plugins
     , ircBotCommands             :: M.Map String [String]   -- Loaded plugins
     , ircBotChan                 :: Chan IrcBotMsg          -- The IrcBot's communication channel
     , ircBotMasterChan           :: Chan BotMsg             -- The Hsbot communication channel
@@ -44,3 +45,6 @@ data IrcServerState = IrcServerState
     , ircServerChan          :: Chan IrcBotMsg -- the IrcBot channel
     }
 
+-- | Utilities for triplets
+first :: (a, b, c) -> a
+first (a, _, _) = a
diff --git a/Hsbot/Plugin.hs b/Hsbot/Plugin.hs
index e545e8d..3feffa8 100644
--- a/Hsbot/Plugin.hs
+++ b/Hsbot/Plugin.hs
@@ -5,8 +5,9 @@ module Hsbot.Plugin
     , unloadPlugin
     ) where
 
-import Control.Concurrent
-import Control.Concurrent.Chan ()
+import Control.Concurrent (forkIO)
+import Control.Concurrent.Chan
+import Control.Concurrent.MVar
 import Control.Exception
 import Control.Monad.State
 import qualified Data.Map as M
@@ -29,12 +30,13 @@ spawnPlugin (IrcBotConfig ircConfig) = do
     bot <- get
     let chan  = botChan bot
     pchan <- liftIO (newChan :: IO (Chan BotMsg))
-    threadId <- liftIO $ forkIO (startIrcbot ircConfig chan pchan)
+    mvar  <- liftIO newEmptyMVar
+    threadId <- liftIO . forkIO $ finally (startIrcbot ircConfig chan pchan) (putMVar mvar ())
     let plugin  = PluginState { pluginName    = ircConfigName ircConfig
                               , pluginChan    = pchan
                               , pluginHandles = M.empty }
         plugins = botPlugins bot
-    put $ bot { botPlugins = M.insert (pluginName plugin) (plugin, threadId) plugins }
+    put $ bot { botPlugins = M.insert (pluginName plugin) (plugin, mvar, threadId) plugins }
     resumeData <- gets botResumeData
     liftIO $ modifyMVar_ resumeData (\oldData -> return $ M.insert (ircConfigName ircConfig) M.empty oldData)
 
@@ -52,9 +54,10 @@ killPlugin name = do
     let oldPlugins = botPlugins bot
     -- We check if the plugin exists
     case M.lookup name oldPlugins of
-        Just (_, threadId) -> do
+        Just (_, mvar, threadId) -> do
             let newPlugins = M.delete name oldPlugins
             liftIO $ throwTo threadId UserInterrupt
             put $ bot { botPlugins = newPlugins }
+            liftIO $ takeMVar mvar
         Nothing            -> return ()
 
diff --git a/Hsbot/Types.hs b/Hsbot/Types.hs
index 009390d..66b4d6b 100644
--- a/Hsbot/Types.hs
+++ b/Hsbot/Types.hs
@@ -27,7 +27,7 @@ type Bot = StateT BotState IO
 -- | An Hsbot state
 data BotState = BotState
     { botStartTime  :: UTCTime      -- the bot's uptime
-    , botPlugins    :: M.Map String (PluginState, ThreadId) -- Loaded plugins
+    , botPlugins    :: M.Map String (PluginState, MVar (), ThreadId) -- Loaded plugins
     , botChan       :: Chan BotMsg  -- the bot's communication channel
     , botConfig     :: [BotConfig]  -- the bot's starting config
     , botResumeData :: MVar BotResumeData   -- the necessary data to resume the bot's operations on reboot