Source

udon / src / Client.hs

{-# LANGUAGE GeneralizedNewtypeDeriving, PatternGuards #-}
module Client where

import Prelude hiding (catch)
import Control.Arrow ((>>>), (&&&), (***))
import Control.Monad (forM_)
import Control.Monad.Reader (ReaderT, runReaderT, ask, asks, MonadReader, MonadIO)
import Data.List (intercalate, mapAccumL)
import Data.Time (UTCTime, getCurrentTime, diffUTCTime)
import Network.Socket (socket, Family(AF_UNIX), SocketType(Stream), defaultProtocol, connect, SockAddr(SockAddrUnix), socketToHandle, Socket)
import System.Environment (getProgName, getArgs)
import System.Exit (exitFailure)
import System.FilePath ((</>), takeBaseName)
import System.IO (Handle, IOMode(ReadWriteMode), hSetBuffering, BufferMode(LineBuffering), hGetLine, hPutStrLn, stderr)
import System.IO.Error (catch)
import System.Posix (usleep)
import System.Posix.Directory (getWorkingDirectory)
import System.Posix.Env (getEnvironment)
import System.Posix.Process (executeFile)
import Text.Printf (printf)

import Paths_udon (version)
import Data.Version (showVersion)

import Common
import Process (parseSignal, runCommand, waitChild')
import Utils (readMaybeIO, hPutStrLnM, putStrLnM, printM, hPrintM, hGetLineM)

type ErrorMessage = String

newtype Client a = Client { runClient :: ReaderT Config IO a }
    deriving (Functor, Monad, MonadIO, MonadReader Config)

execMaster :: IO ()
execMaster = do
    masterPath <- getMasterPath
    env <- getEnvironment
    cwd <- getWorkingDirectory
    (_, _, pid) <- runCommand masterPath [] cwd env
    waitChild' pid
    return ()

startClient :: Config -> [String] -> IO ()
startClient cfg args = runReaderT (runClient $ client args) cfg

client :: [String] -> Client ()
client []                       = printHelp
client ("help":args)            = printHelp
client ("ping":[])              = runCCommand commandPing
client ("register":args)        = runCCommand (commandRegister args)
client ("unregister":daemon:[]) = runCCommand (commandUnregister daemon)
client ("start":daemon:[])      = runCCommand (commandStart daemon)
client ("stop":daemon:[])       = runCCommand (commandStop daemon "SIGTERM")
client ("stop":daemon:sig:[])   = runCCommand (commandStop daemon sig)
client ("kill":daemon:[])       = runCCommand (commandKill daemon "SIGTERM")
client ("kill":daemon:sig:[])   = runCCommand (commandKill daemon sig)
client ("list":["-v"])          = runCCommand (commandList ["-v"])
client ("list":[])              = runCCommand (commandList [])
client ("info":"-v":daemon:[])  = runCCommand (commandInfoV daemon)
client ("info":daemon:[])       = runCCommand (commandInfo daemon)
client ("tail":args@(_:_) )     = runCCommand (commandTail args)
client ("less":args@(_:_) )     = runCCommand (commandLess args)
client ("shutdown-master":[])   = runCCommand commandShutdownMaster
client ("restart-master":[])    = runCCommand commandRestartMaster
client ("version":[])           = runCCommand commandVersion
client ("master-version":[])    = runCCommand commandMasterVersion
client _                        = printHelp >> io exitFailure

type ClientCommand = Client (Either ErrorMessage String)

runCCommand :: ClientCommand -> Client ()
runCCommand c = either (\e -> hPutStrLnM stderr e >> io exitFailure) putStrLnM =<< c

printHelp :: Client ()
printHelp = do
    self <- io getProgName
    io $ mapM_ putStrLn
        [ "Usage: " ++ self ++ " <command> [args...]"
        , "commands:"
        , "    help:                                     print this message"
        , "    ping:                                     pong"
        , "    register [-n <name>] <cmd> [cmdargs...]:  register <cmd> as daemon"
        , "    unregister <daemon>:                      unregister daemon"
        , "    start <daemon>:                           start <daemon>"
        , "    stop <daemon> [<signal>]:                 stop <daemon> with <signal>"
        , "    kill <daemon> [<signal>]:                 send <signal> to <daemon>"
        , "    list [-v]:                                get list of daemons"
        , "                                               -v  show loggers"
        , "    info [-v] <daemon>:                       get daemon's info"
        , "                                               -v  show more detail info with logger's info"
        , "    tail [tail's options...] <daemon>:        tail daemon's log"
        , "    less [less's options...] <daemon>:        less daemon's log"
        , "    shutdown-master                           shutdown master process"
        , "    restart-master                            restart/upgrade master process while daemon running"
        , "    verion:                                   show version"
        , "    master-verion:                            show master's version"
        ]

commandPing :: ClientCommand
commandPing = do
    h <- connectMaster
    hPrintM h Ping
    rsp <- hGetResponse h
    case rsp of
        Nothing     -> return $ Left "Unknown Response"
        Just (NG e) -> return $ Left ("NG: " ++ e)
        Just Pong   -> return $ Right "Pong"
        Just r      -> return $ Left ("Invalid Response: " ++ show r)

commandRegister :: [String] -> ClientCommand
commandRegister ("-n":name:cmd:args)      = commandRegister' name cmd args
commandRegister (('-':'n':name):cmd:args) = commandRegister' name cmd args
commandRegister (cmd:args)                = commandRegister' (takeBaseName cmd) cmd args

commandRegister' :: String -> String -> [String] -> ClientCommand
commandRegister' name cmd args = do
    h <- connectMaster
    cwd <- io getWorkingDirectory
    env <- io getEnvironment
    let d = Daemon
                { dName      = name
                , dCommand   = cmd
                , dArgs      = args
                , dCwd       = cwd
                , dEnv       = env
                , dAutoStart = False
                }
    hPrintM h $ RegisterDaemon d
    assertResponseOK =<< hGetResponse h

commandUnregister :: String -> ClientCommand
commandUnregister dn = do
    h <- connectMaster
    hPrintM h $ UnregisterDaemon dn
    assertResponseOK =<< hGetResponse h

commandStart :: String -> ClientCommand
commandStart dn = do
    h <- connectMaster
    hPrintM h $ StartDaemon dn
    assertResponseOK =<< hGetResponse h

commandStop :: String -> String -> ClientCommand
commandStop dn sig = do
    h <- connectMaster
    hPrintM h $ SetAutoStart dn False
    ret <- assertResponseOK =<< hGetResponse h
    case ret of
        e@(Left _) -> return e
        Right _    -> commandKill dn sig

commandKill :: String -> String -> ClientCommand
commandKill dn sig = do
    h <- connectMaster
    let msig = parseSignal sig
    case msig of
        Nothing    -> return $ Left ("Invalid signal name: " ++ sig)
        Just sig' -> do
            hPrintM h $ KillDaemon dn sig'
            assertResponseOK =<< hGetResponse h

commandList :: [String] -> ClientCommand
commandList ("-v":_) = do
    now <- io getCurrentTime
    commandListWith (printDaemonStatus now *** printLoggerStatus now >>> \(a, b) -> [a, b])
commandList _ = do
    now <- io getCurrentTime
    commandListWith (fst >>> printDaemonStatus now >>> (:[]))

commandListWith :: ((DaemonState, LoggerState) -> [(String, String)]) -> ClientCommand
commandListWith printer = do
    h <- connectMaster
    hPrintM h GetListDaemons
    rsp <- hGetResponse h
    case rsp of
        Nothing               -> return $ Left "Unknown Response"
        Just (NG e)           -> return $ Left ("NG: " ++ e)
        Just (ListDaemons dls) -> return $ Right $ showAList 20 $ ("DAEMON", "STATUS") : concatMap printer dls
        Just r                -> return $ Left ("Invalid Response: " ++ show r)

printDaemonStatus :: UTCTime -> (DaemonState -> (String, String))
printDaemonStatus now = dName . dsDaemon &&& showDaemonStatus now . dsStatus

printLoggerStatus :: UTCTime -> (LoggerState -> (String, String))
printLoggerStatus now = printDaemonStatus now >>> ("  (" ++) *** (++ ")")

showAList :: Int -> [(String, String)] -> String
showAList n al = unlines $ map showKV al
  where
    showKV (k, v) = pad n k ++ " " ++ v
    pad i cs
        | length cs > i = cs ++ '\n' : replicate i ' '
        | otherwise     = cs ++ replicate (i - length cs) ' '

commandInfo :: String -> ClientCommand
commandInfo dn = commandInfoWith dn $ \dst _ -> simplePprintDaemonState dst

commandInfoV :: String -> ClientCommand
commandInfoV dn = commandInfoWith dn $ \dst lst -> pprintDaemonState dst ++ pprintDaemonState lst

commandInfoWith :: String -> (DaemonState -> LoggerState -> String) -> ClientCommand
commandInfoWith dn printer = do
    h <- connectMaster
    hPrintM h $ GetDaemonStatus dn
    rsp <- hGetResponse h
    case rsp of
        Nothing                     -> return $ Left "Unknown Response"
        Just (NG e)                 -> return $ Left ("NG: " ++ e)
        Just (DaemonStatus dst lst) -> return $ Right $ printer dst lst
        Just r                      -> return $ Left ("Invalid Response: " ++ show r)

commandTail :: [String] -> ClientCommand
commandTail args = commandWithCommand "tail" args

commandLess :: [String] -> ClientCommand
commandLess args = commandWithCommand "less" args

commandWithCommand :: String -> [String] -> ClientCommand
commandWithCommand cmd args = do
    logd <- asks logDir
    let logf  = logd </> (escapeFName $ dn ++ ".log")
        dn    = last args
        args' = init args ++ [logf]
    io $ executeFile cmd True args' Nothing
    return . Left $ "Cannot execute command: " ++ show (cmd:args')

commandShutdownMaster :: ClientCommand
commandShutdownMaster = do
    h <- connectMaster
    hPrintM h ShutdownMaster
    assertResponseOK =<< hGetResponse h

commandRestartMaster :: ClientCommand
commandRestartMaster = do
    h <- connectMaster
    masterPath <- io getMasterPath
    hPrintM h $ RestartMaster masterPath ["--resume"]
    assertResponseOK =<< hGetResponse h

commandVersion :: ClientCommand
commandVersion = do
    self <- io getProgName
    return $ Right (self ++ " " ++ showVersion version)

commandMasterVersion :: ClientCommand
commandMasterVersion = do
    h <- connectMaster
    hPrintM h $ GetMasterVersion
    rsp <- hGetResponse h
    case rsp of
        Nothing                  -> return $ Left "Unknown Response"
        Just (NG e)              -> return $ Left ("NG: " ++ e)
        Just (MasterVersion n v) -> return $ Right (n ++ " " ++ showVersion v)
        Just r                   -> return $ Left ("Invalid Response: " ++ show r)

connectMaster :: Client Handle
connectMaster = do
    sp <- asks masterSock
    io $ do
        s <- socket AF_UNIX Stream defaultProtocol
        tryConnect s (SockAddrUnix sp)
        h <- socketToHandle s ReadWriteMode
        hSetBuffering h LineBuffering
        return h
  where
    tryConnect s addr = tryConnect' 0 10 s addr

    tryConnect' count limit s addr
        | count > limit = fail $ "Cannot connect to " ++ show addr
        | otherwise     = connect s addr `catch` \_ -> usleep 100000 >> tryConnect' (count + 1) limit s addr

hGetResponse :: Handle -> Client (Maybe Response)
hGetResponse h = io . readMaybeIO =<< hGetLineM h

assertResponseOK :: Maybe Response -> ClientCommand
assertResponseOK Nothing       = return $ Left "Unknown Response"
assertResponseOK (Just (NG e)) = return $ Left ("NG: " ++ e)
assertResponseOK (Just OK)     = return $ Right "OK"
assertResponseOK (Just r)      = return $ Left ("Invalid Response: " ++ show r)

showDaemonStatus :: UTCTime -> DaemonStatus -> String
showDaemonStatus now (Running pid t)  = "running (pid: " ++ show pid ++ ") " ++ showSeconds (truncate $ diffUTCTime now t)
showDaemonStatus _   (Exited code)    = "exited (eixt with: " ++ show code ++ ")"
showDaemonStatus _   (Terminated sig) = "terminated (by signal " ++ show sig ++ ")"
showDaemonStatus _   Initializing     = "initializing"
showDaemonStatus _   NotStarted       = "not started"
showDaemonStatus _   Deleting         = "deleting"
showDaemonStatus now (FastSpawn ds)   = "respawning too fast, stopped (" ++ showDaemonStatus now ds ++ ")"

showSeconds :: Integer -> String
showSeconds sec
    | (0, h, m, s) <- secToDHMS sec = printf "%02d:%02d:%02d" h m s
    | (1, h, m, s) <- secToDHMS sec = printf "1 day %02d:%02d" h m
    | (d, h, m, s) <- secToDHMS sec = printf "%d days %02d:%02d" d h m

secToDHMS :: Integer -> (Integer, Integer, Integer, Integer)
secToDHMS i = let (d, [s, m, h]) = mapAccumL divMod i [60, 60, 24]
              in (d, h, m, s)