bein-0.1: Bein/Minion.hs

-- Add setting status to running in db, and checking state and finalizing status when doing cleanup
module Main where

import Database.HDBC
import System.Posix.Syslog
import System.Environment (getProgName)
import Control.Monad.Reader
import Database.HDBC.PostgreSQL
import Control.Concurrent.STM
import System.Exit (ExitCode(..))
import System.Directory
import System.Posix.Signals
import System.Posix.Process
import System.FilePath.Posix (joinPath)

import Bein.Configuration
import Bein.Minion.Types
import Bein.Minion.Arguments
import Bein.SocketHandler
import Bein.Minion.Commands
import Bein.Minion.Protocol

main :: IO ()
main = do
  getSettings >>= getState >>= \st -> do
    programName <- getProgName
    withSyslog programName [PID,PERROR] DAEMON $ runReaderT main' st
      
getState :: Settings -> IO State
getState opts = do
  conn <- connectPostgreSQL "dbname=bein"
  config <- readConfiguration conn
  configTV <- newTVarIO config
  minP <- newTVarIO Nothing
  return $ State { stDb = conn, stConfigT = configTV, stSettings = opts, stMinionPid = minP }

main' :: ReaderT State IO b
main' = do 
  execution >>= \e -> checkExecution e >>= \r -> case r of
    Left errors -> dieWith errors
    Right _ -> do installHandlers
                  scratchDir >>= createScratchDirectory
                  script <- execution >>= \_ -> scratchDir >>= \s -> writeScript e s
                  forkR $ minionSocket >>= listenWith minionProtocol
                  runScript script
                  forever $ return () -- Doing this requires using exitImmediately to kill the program when we're done.

checkExecution :: ExecutionID -> BeinM State (Either String ())
checkExecution ex = do
  [notFuture, isWaiting, nonemptyScript, noNullInputs] <- 
    liftM (map fromSql) $ receiveOneRow $ query queryString [toSql ex]
  let errors = concat $ map snd $ filter (not.fst) $
               [(notFuture, "Execution " ++ show ex ++ " is not a future.\n"),
                (isWaiting, "Execution " ++ show ex ++ " is not waiting.\n"),
                (nonemptyScript, "Execution " ++ show ex ++ 
                                 " refers to a program with an empty script.\n"),
                (noNullInputs, "Execution " ++ show ex ++ " has null inputs.\n")]
    in if (notFuture && isWaiting && nonemptyScript && noNullInputs) 
         then do verboseMsgLn "Setting status to running."
                 update "update executions set status = 'running' where id=?" [toSql ex]
                 return (Right ())
         else return (Left errors)
    where queryString = "select not_future, is_waiting, nonempty_script, " ++ 
                        "no_null_inputs from runnable where id = ?"

installHandlers :: BeinM State ()
installHandlers = ask >>= \s -> do 
  verboseMsg "Installing handlers..."
  mapM_ ignoreSignal [sigALRM, sigPIPE, sigPOLL, sigPROF, sigUSR1, sigUSR2, sigVTALRM]
  liftIO $ installHandler sigCHLD (Catch $ runReaderT sigCHLDHandler s) (Just fullSignalSet)
  liftIO $ installHandler sigHUP (Catch $ runReaderT (cleanUp (ExitFailure 1)) s) (Just fullSignalSet)
  liftIO $ installHandler sigTERM (Catch $ runReaderT (cleanUp (ExitFailure 1)) s) (Just fullSignalSet)
  verboseMsgLn "ok"

sigCHLDHandler :: BeinM State ()
sigCHLDHandler = do
  v <- liftIO $ getAnyProcessStatus False False
  case v of
    Just (pid,status) -> do 
      minionPid >>= \p -> if p /= Just pid 
        then verboseMsgLn "Received sigCHLD from something besides script." >> sigCHLDHandler
        else (verboseMsgLn $ "Child exited with status " ++ show status) >> cleanUp ExitSuccess
    Nothing -> return ()
                                                  

createScratchDirectory :: FilePath -> BeinM State ()
createScratchDirectory dir = do
  verboseMsg $ "Creating scratch directory " ++ dir ++ "..."
  (liftIO $ createDirectory dir) `catchR` 
    (\e -> dieWith $ "Could not create scratch directory.\n" ++ show e)
  (liftIO $ setCurrentDirectory dir) `catchR` 
    (\e -> dieWith $ "Could not set working directory.\n" ++ show e)
  verboseMsgLn "ok"

writeScript :: ExecutionID -> FilePath -> BeinM State FilePath
writeScript ex dir = do 
  verboseMsg $ "Writing script to " ++ dir ++ "..."
  [lang,script] <- liftM (map fromSql) $ receiveOneRow $ query queryString [toSql ex]
  let scriptName = case lang of
        "perl" -> "script.pl"
        "r" -> "script.R"
        _ -> error "Invalid value for language."
  scriptPath <- scratchDir >>= (\d -> return $ joinPath [d,scriptName])
  liftIO $ writeFile scriptPath script
  verboseMsgLn "ok"
  return scriptPath
    where queryString = "select p.language,p.script from programs " ++ 
                        "as p join executions as e on e.program=p.id where e.id=?"
                     
runScript :: FilePath -> BeinM State ()
runScript script = do
  verboseMsg $ "Running script " ++ script ++ "..."
  perlCmd <- configField perl_executable
  minionPort <- configField minion_port
  p <- liftIO $ forkProcess $ executeFile perlCmd False [script,minionPort] Nothing
  writeMinionPid p
  verboseMsgLn "ok"