tpar-0.1.0.0: Main.hs
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE OverloadedStrings #-}
import Control.Monad.Catch
import Control.Monad (when, unless, forever, replicateM_)
import Control.Monad.IO.Class
import Control.Error hiding (err)
import Data.Foldable
import qualified Data.Map.Strict as M
import Data.Time.Clock
import Data.Time.Format.Human
import System.Exit
import System.IO (stderr, stdout)
import Data.Aeson as Aeson
import qualified Data.ByteString.Lazy.Char8 as BS.L
import qualified Text.PrettyPrint.ANSI.Leijen as T.PP
import Text.PrettyPrint.ANSI.Leijen (Doc, (<+>), (<$$>))
import qualified Network.Transport.TCP as TCP
import Network.Socket (ServiceName, HostName)
import Network.BSD (getHostName)
import Options.Applicative hiding (action, header)
import qualified Options.Applicative as O
import Control.Concurrent (threadDelay)
import Control.Distributed.Process
import Control.Distributed.Process.Internal.Types (NodeId(..))
import Control.Distributed.Process.Node
import qualified Text.Trifecta as TT
import Pipes
import qualified Pipes.Concurrent as P.C
import qualified Pipes.Prelude as P.P
import TPar.Rpc
import TPar.ProcessPipe hiding (runProcess)
import TPar.Server
import TPar.SubPubStream as SubPub
import TPar.Server.Types
import TPar.JobMatch
import TPar.Types
portOption :: Mod OptionFields ServiceName -> Parser ServiceName
portOption m =
option str
( short 'p' <> long "port"
<> value "5757" <> m
)
hostOption :: Parser HostName
hostOption =
strOption ( short 'H' <> long "host" <> value "localhost" <> help "server host name" )
jobMatchArg :: Parser JobMatch
jobMatchArg = argument (liftTrifecta parseJobMatch) (help "job match expression")
type Mode = IO ()
tpar :: ParserInfo Mode
tpar = O.info (helper <*> tparParser)
$ O.fullDesc
<> O.progDesc "Start queues, add workers, and enqueue tasks"
<> O.header "tpar - simple distributed task queuing"
tparParser :: Parser Mode
tparParser =
subparser
$ command "server" ( info modeServer
$ fullDesc <> progDesc "Start a server")
<> command "worker" ( info modeWorker
$ fullDesc <> progDesc "Start a worker")
<> command "enqueue" ( info modeEnqueue
$ fullDesc <> progDesc "Enqueue a job")
<> command "status" ( info modeStatus
$ fullDesc <> progDesc "Show queue status")
<> command "kill" ( info modeKill
$ fullDesc <> progDesc "Kill or dequeue a job")
<> command "rerun" ( info modeRerun
$ fullDesc <> progDesc "Restart a failed job")
<> command "watch" ( info modeWatch
$ fullDesc <> progDesc "Watch the output of a set of jobs")
<> command "dump" ( info modeDump
$ fullDesc <> progDesc "Dump details about a set of jobs in JSON")
withServer' :: HostName -> ServiceName
-> (ServerIface -> Process a) -> Process a
withServer' host port action = do
let nid :: NodeId
nid = NodeId (TCP.encodeEndPointAddress host port 0)
-- request server interface
mref <- monitorNode nid
(sq, rq) <- newChan :: Process (SendPort ServerIface, ReceivePort ServerIface)
nsendRemote nid "tpar" sq
iface <- receiveWait
[ matchIf (\(NodeMonitorNotification ref _ _) -> ref == mref) $
\(NodeMonitorNotification _ _ reason) ->
case reason of
DiedDisconnect -> fail "Failed to connect. Are you sure there is a server running?"
_ -> fail $ show reason
, matchChan rq pure
]
unless (protocolVersion iface == currentProtocolVersion) $
fail $ "tpar protocol version mismatch: server speaks "++show (protocolVersion iface)
++ ", client speaks "++show currentProtocolVersion
-- request server interface
link (serverPid iface)
unlinkNode nid
r <- action iface
unlink (serverPid iface)
return r
withServer :: HostName -> ServiceName
-> (ServerIface -> Process ()) -> IO ()
withServer host port action = do
hostname <- getHostName
Right transport <- TCP.createTransport hostname "0" TCP.defaultTCPParameters
node <- newLocalNode transport initRemoteTable
runProcess node $ withServer' host port action
modeWorker :: Parser Mode
modeWorker =
run <$> option (auto >>= checkNWorkers)
( short 'N' <> long "workers" <> value 1
<> help "number of local workers to start"
)
<*> hostOption
<*> portOption (help "server port number")
<*> option (Just <$> (auto <|> pure 10))
( short 'r' <> long "reconnect" <> metavar "SECONDS" <> value Nothing
<> help "attempt to reconnect when server vanishes (with optional retry period); otherwise terminates on server vanishing"
)
<* helper
where
checkNWorkers n
| n >= 1 = return n
| otherwise = fail "Worker count (-N) should be at least one"
run nWorkers serverHost serverPort reconnectPeriod =
perhapsRepeat $ withServer serverHost serverPort $ \serverIface -> do
replicateM_ nWorkers $ spawnLocal $ runRemoteWorker serverIface
liftIO $ forever threadDelay maxBound
where
perhapsRepeat action
| Just period <- reconnectPeriod = forever $ do
handleAll (liftIO . print) action
liftIO (threadDelay $ 1000*1000*period)
| otherwise = action
modeServer :: Parser Mode
modeServer =
run <$> option str (short 'H' <> long "host" <> value "localhost"
<> help "interface address to listen on" )
<*> portOption (help "port to listen on")
<*> option auto ( short 'N' <> long "workers" <> value 0
<> help "number of local workers to start"
)
<* helper
where
run serverHost serverPort nLocalWorkers = do
Right transport <- TCP.createTransport serverHost serverPort TCP.defaultTCPParameters
node <- newLocalNode transport initRemoteTable
runProcess node $ do
iface <- runServer
replicateM_ nLocalWorkers $ spawnLocal $ runRemoteWorker iface
liftIO $ forever $ threadDelay maxBound
modeEnqueue :: Parser Mode
modeEnqueue =
run <$> hostOption
<*> portOption (help "server port number")
<*> sinkType
<*> switch (short 'w' <> long "watch" <> help "Watch output of task")
<*> option (JobName <$> str)
(short 'n' <> long "name" <> value (JobName "unnamed-job")
<> help "Set the job's name")
<*> option str
(short 'd' <> long "directory" <> value "."
<> help "Set the directory the job will be launched from (relative to the working directory of the worker who runs it)")
<*> option (Priority <$> auto)
(short 'P' <> long "priority" <> value (Priority 0)
<> help "Set the job's priority")
<*> some (argument str idm)
<* helper
where
run :: HostName -> ServiceName
-> OutputStreams (Maybe FilePath) -> Bool
-> JobName -> FilePath -> Priority -> [String]
-> IO ()
run serverHost serverPort sink watch name dir priority (cmd:args) =
withServer serverHost serverPort $ \iface -> do
let jobReq = JobRequest { jobName = name
, jobPriority = priority
, jobCommand = cmd
, jobArgs = args
, jobSinks = sink
, jobCwd = dir
, jobEnv = Nothing
}
if watch
then do
prod <- enqueueAndFollow iface jobReq
code <- runEffect $ prod >-> P.P.mapM_ (processOutputToHandles $ OutputStreams stdout stderr)
case code of
ExitSuccess -> return ()
ExitFailure n ->
liftIO $ putStrLn $ "exited with code "++show n
else do
Right _ <- callRpc (enqueueJob iface) (jobReq, Nothing)
return ()
run _ _ _ _ _ _ _ _ = fail "Expected command line"
sinkType :: Parser (OutputStreams (Maybe FilePath))
sinkType =
OutputStreams
<$> option (Just <$> str)
(short 'o' <> long "output" <> metavar "FILE"
<> help "remote file to log standard output to"
<> value Nothing)
<*> option (Just <$> str)
(short 'e' <> long "error" <> metavar "FILE"
<> help "remote file to log standard error to"
<> value Nothing)
liftTrifecta :: TT.Parser a -> ReadM a
liftTrifecta parser = do
s <- str
case TT.parseString parser mempty s of
TT.Success a -> return a
TT.Failure err -> fail $ show $ TT._errDoc err
modeWatch :: Parser Mode
modeWatch =
run <$> hostOption
<*> portOption (help "server port number")
<*> jobMatchArg
<* helper
where
run serverHost serverPort jobMatch =
withServer serverHost serverPort $ \iface -> do
inputs <- subscribeToJobs iface jobMatch
let input = foldMap fold inputs
failed = M.keys $ M.filter isNothing inputs
unless (null failed) $ liftIO $ print
$ T.PP.red "warning: failed to attach to"
<+> prettyShow (length failed) <+> "jobs"
<$$> T.PP.nest 4 (T.PP.vcat $ map prettyShow failed)
runEffect $ P.C.fromInput input
>-> P.P.mapM_ (processOutputToHandles $ OutputStreams stdout stderr)
subscribeToJobs :: ServerIface -> JobMatch
-> Process (M.Map JobId (Maybe (P.C.Input ProcessOutput)))
subscribeToJobs iface jobMatch = do
Right jobs <- callRpc (getQueueStatus iface) jobMatch
prods <- traverse SubPub.subscribe
$ M.unions
[ M.singleton jobId jobMonitor
| Job {..} <- jobs
, Running {..} <- pure jobState
]
traverse (traverse producerToInput) prods
where
producerToInput :: Producer ProcessOutput Process ExitCode
-> Process (P.C.Input ProcessOutput)
producerToInput prod = do
(output, input) <- liftIO $ P.C.spawn (P.C.bounded 1)
void $ spawnLocal $ runEffect $ void prod >-> P.C.toOutput output
return input
modeStatus :: Parser Mode
modeStatus =
run <$> hostOption
<*> portOption (help "server port number")
<*> switch (short 'v' <> long "verbose" <> help "verbose queue status")
<*> (jobMatchArg <|> pure (NegMatch NoMatch))
<* helper
where
run serverHost serverPort verbose jobMatch =
withServer serverHost serverPort $ \iface -> do
Right jobs <- callRpc (getQueueStatus iface) jobMatch
time <- liftIO getCurrentTime
let prettyTime = T.PP.text . humanReadableTime' time
liftIO $ T.PP.putDoc $ T.PP.vcat $ map (prettyJob verbose prettyTime) jobs ++ [mempty]
modeDump :: Parser Mode
modeDump =
run <$> hostOption
<*> portOption (help "server port number")
<*> (jobMatchArg <|> pure (NegMatch NoMatch))
<* helper
where
run serverHost serverPort jobMatch =
withServer serverHost serverPort $ \iface -> do
Right jobs <- callRpc (getQueueStatus iface) jobMatch
liftIO $ BS.L.putStrLn $ Aeson.encode $ map jobToJson jobs
jobToJson :: Job -> Value
jobToJson (Job {jobRequest = JobRequest{..}, ..}) =
object [ "id" .= fromEnum jobId
, "name" .= case jobName of JobName s -> s
, "priority" .= fromEnum jobPriority
, "command" .= jobCommand
, "arguments" .= jobArgs
, "working_dir" .= jobCwd
, "remote_logs" .= object [ "stdout" .= stdOut jobSinks
, "stderr" .= stdErr jobSinks
]
, "environment" .= jobEnv
, "state" .= state
]
where
stateObj s rest = object $ ("state" .= (s :: String)) : rest
state =
case jobState of
Queued {..} ->
stateObj "queued"
[ "queued_at" .= jobQueueTime ]
Starting {..} ->
stateObj "starting"
[ "queued_at" .= jobQueueTime
, "starting_since" .= jobStartingTime
, "running_on" .= show jobProcessId
]
Running {..} ->
stateObj "running"
[ "queued_at" .= jobQueueTime
, "started_at" .= jobStartTime
, "running_on" .= show jobProcessId
]
Finished {..} ->
stateObj "finished"
[ "queued_at" .= jobQueueTime
, "started_at" .= jobStartTime
, "finished_at" .= jobFinishTime
, "ran_on" .= show jobWorkerNode
, "exit_code" .= getExitCode jobExitCode
]
Failed {..} ->
stateObj "finished"
[ "queued_at" .= jobQueueTime
, "started_at" .= jobStartTime
, "failed_at" .= jobFailedTime
, "ran_on" .= show jobWorkerNode
, "error" .= jobErrorMsg
]
Killed {..} ->
stateObj "killed"
[ "queued_at" .= jobQueueTime
, "started_at" .= jobKilledStartTime
, "killed_at" .= jobKilledTime
, "ran_on" .= show jobKilledWorkerNode
]
prettyJob :: Bool -> (UTCTime -> Doc) -> Job -> Doc
prettyJob verbose prettyTime (Job {..}) =
T.PP.vcat $ [header] ++ (if verbose then [details] else [])
where
JobRequest {..} = jobRequest
twoCols :: Int -> [(Doc, Doc)] -> Doc
twoCols width =
T.PP.vcat . map (\(a,b) -> T.PP.fillBreak width a <+> T.PP.align b)
header =
T.PP.fillBreak 5 (prettyJobId jobId)
<+> T.PP.fillBreak 50 (prettyJobName jobName)
<+> prettyJobState jobState
details =
T.PP.indent 4 $ twoCols 15
[ ("priority:", prettyPriority jobPriority)
, ("queued:", prettyTime (jobQueueTime jobState))
, ("command:", T.PP.text jobCommand)
, ("arguments:", T.PP.hsep $ map T.PP.text jobArgs)
, ("logging:", T.PP.vcat $ toList $ (\x y -> x<>":"<+>y)
<$> OutputStreams "stdout" "stderr"
<*> fmap (maybe "none" T.PP.text) jobSinks)
, ("status:", prettyDetailedState jobState)
]
<$$> mempty
prettyDetailedState Queued{..} =
"waiting to run" <+> T.PP.parens ("since" <+> prettyTime jobQueueTime)
prettyDetailedState Starting{..} =
"starting on" <+> prettyShow jobProcessId
<+> T.PP.parens ("since" <+> prettyTime jobStartingTime)
prettyDetailedState Running{..} =
"running on" <+> prettyShow jobProcessId
<+> T.PP.parens ("since" <+> prettyTime jobStartTime)
prettyDetailedState Finished{..} =
"finished with" <+> prettyShow (getExitCode jobExitCode)
<+> T.PP.parens ("at" <+> prettyTime jobFinishTime)
<$$> "started at" <+> prettyTime jobStartTime
<$$> "ran on" <+> prettyShow jobWorkerNode
prettyDetailedState Failed{..} =
"failed with error" <+> T.PP.parens (prettyTime jobFailedTime)
<$$> "started at" <+> prettyTime jobStartTime
<$$> T.PP.indent 4 (T.PP.text jobErrorMsg)
prettyDetailedState Killed{..} =
"killed at user request" <+> T.PP.parens (prettyTime jobKilledTime)
<$$> maybe "never started" (\t -> "started at" <+> prettyTime t) jobKilledStartTime
prettyJobState Queued{} = T.PP.blue "queued"
prettyJobState Starting{} = T.PP.blue "starting"
prettyJobState Running{} = T.PP.cyan "running"
prettyJobState Finished{..}
| ExitFailure _ <- jobExitCode = T.PP.yellow $ "failed"
| otherwise = T.PP.green "finished"
prettyJobState Failed{} = T.PP.red "job failed"
prettyJobState Killed{} = T.PP.yellow "killed"
prettyJobId (JobId n) = T.PP.int n
prettyJobName (JobName name) = T.PP.bold $ T.PP.text name
prettyPriority (Priority p) = T.PP.int p
getExitCode :: ExitCode -> Int
getExitCode ExitSuccess = 0
getExitCode (ExitFailure c) = c
prettyShow :: Show a => a -> Doc
prettyShow = T.PP.text . show
modeKill :: Parser Mode
modeKill =
run <$> hostOption
<*> portOption (help "server port number")
<*> jobMatchArg
<* helper
where
run serverHost serverPort jobMatch =
withServer serverHost serverPort $ \iface -> do
Right jobs <- callRpc (killJobs iface) jobMatch
liftIO $ T.PP.putDoc $ T.PP.vcat $ map (prettyJob False prettyShow) jobs ++ [mempty]
liftIO $ when (null jobs) $ exitWith $ ExitFailure 1
modeRerun :: Parser Mode
modeRerun =
run <$> hostOption
<*> portOption (help "server port number")
<*> jobMatchArg
<* helper
where
run serverHost serverPort jobMatch =
withServer serverHost serverPort $ \iface -> do
Right jobs <- callRpc (rerunJobs iface) jobMatch
liftIO $ T.PP.putDoc $ T.PP.vcat $ map (prettyJob False prettyShow) jobs ++ [mempty]
liftIO $ when (null jobs) $ exitWith $ ExitFailure 1
main :: IO ()
main = do
run <- execParser tpar
res <- runExceptT $ tryIO run
case res of
Left err -> putStrLn $ "error: "++show err
Right () -> return ()