{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Main
( main
) where
import Control.Monad (void, when)
import Control.Monad.IO.Class (liftIO)
import Data.Binary (decodeFile, encodeFile)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B (lines, pack, putStr, readFile,
split, unpack)
import Data.Int (Int64)
import Data.List (isPrefixOf, transpose)
import Data.Maybe (fromMaybe)
import Metro.Class (Transport)
import Metro.Socket (getHost, getService)
import Metro.TP.Socket (socket)
import Metro.TP.TLS (makeClientParams', tlsConfig)
import Metro.TP.WebSockets (clientConfig)
import Metro.TP.XOR (xorConfig)
import Periodic.Trans.Client
import Periodic.Types (Workload (..))
import System.Environment (getArgs, lookupEnv)
import System.Exit (exitSuccess)
import Text.Read (readMaybe)
import Data.String (fromString)
import Data.UnixTime
import System.IO.Unsafe (unsafePerformIO)
import qualified Text.PrettyPrint.Boxes as T
data Command = Status
| Submit
| Run
| Remove
| Drop
| Config
| Shutdown
| Dump
| Load
| Help
deriving (Eq)
parseCommand :: String -> Command
parseCommand "status" = Status
parseCommand "submit" = Submit
parseCommand "run" = Run
parseCommand "remove" = Remove
parseCommand "drop" = Drop
parseCommand "config" = Config
parseCommand "shutdown" = Shutdown
parseCommand "dump" = Dump
parseCommand "load" = Load
parseCommand _ = Help
data Options = Options
{ host :: String
, xorFile :: FilePath
, useTls :: Bool
, useWs :: Bool
, hostName :: String
, certKey :: FilePath
, cert :: FilePath
, caStore :: FilePath
}
options :: Maybe String -> Maybe String -> Options
options h f = Options { host = fromMaybe "unix:///tmp/periodic.sock" h
, xorFile = fromMaybe "" f
, useTls = False
, useWs = False
, hostName = "localhost"
, certKey = "client-key.pem"
, cert = "client.pem"
, caStore = "ca.pem"
}
parseOptions :: [String] -> Options -> (Command, Options, [String])
parseOptions [] opt = (Help, opt, [])
parseOptions ("-H":x:xs) opt = parseOptions xs opt { host = x }
parseOptions ("--host":x:xs) opt = parseOptions xs opt { host = x }
parseOptions ("--xor":x:xs) opt = parseOptions xs opt { xorFile = x }
parseOptions ("--tls":xs) opt = parseOptions xs opt { useTls = True }
parseOptions ("--ws":xs) opt = parseOptions xs opt { useWs = True }
parseOptions ("--hostname":x:xs) opt = parseOptions xs opt { hostName = x }
parseOptions ("--cert-key":x:xs) opt = parseOptions xs opt { certKey = x }
parseOptions ("--cert":x:xs) opt = parseOptions xs opt { cert = x }
parseOptions ("--ca":x:xs) opt = parseOptions xs opt { caStore = x }
parseOptions (x:xs) opt = (parseCommand x, opt, xs)
printHelp :: IO ()
printHelp = do
putStrLn "periodic - Periodic task system client"
putStrLn ""
putStrLn "Usage: periodic [--host|-H HOST] [--xor FILE|--ws|--tls [--hostname HOSTNAME] [--cert-key FILE] [--cert FILE] [--ca FILE]] command [options]"
putStrLn ""
putStrLn "Commands:"
putStrLn " status Show status"
putStrLn " submit Submit job"
putStrLn " run Run job and output result"
putStrLn " remove Remove job"
putStrLn " drop Drop function"
putStrLn " config Set or Get config"
putStrLn " shutdown Shutdown periodicd"
putStrLn " dump Dump jobs"
putStrLn " load Load jobs"
putStrLn " help Shows a list of commands or help for one command"
putStrLn ""
putStrLn "Available options:"
putStrLn " -H --host Socket path [$PERIODIC_PORT]"
putStrLn " eg: tcp://:5000 (optional: unix:///tmp/periodic.sock) "
putStrLn " --xor XOR Transport encode file [$XOR_FILE]"
putStrLn " --tls Use tls transport"
putStrLn " --ws Use websockets transport"
putStrLn " --hostname Host name"
putStrLn " --cert-key Private key associated"
putStrLn " --cert Public certificate (X.509 format)"
putStrLn " --ca trusted certificates"
putStrLn ""
putStrLn "Version: v1.1.7.1"
putStrLn ""
exitSuccess
printWorkloadHelp :: IO ()
printWorkloadHelp = do
putStrLn ""
putStrLn "Available options:"
putStrLn " -w --workload WORKLOAD or @FILE"
printSubmitHelp :: IO ()
printSubmitHelp = do
putStrLn "periodic submit - Submit job"
putStrLn ""
putStrLn "Usage: periodic submit funcname jobname [-w|--workload WORKLOAD|@FILE] [--later 0]"
printWorkloadHelp
putStrLn " --later Sched job later"
putStrLn ""
exitSuccess
printRunHelp :: IO ()
printRunHelp = do
putStrLn "periodic run - Run job and output result"
putStrLn ""
putStrLn "Usage: periodic run funcname jobname [-w|--workload WORKLOAD|@FILE]"
printWorkloadHelp
putStrLn ""
exitSuccess
printRemoveHelp :: IO ()
printRemoveHelp = do
putStrLn "periodic remove - Remove job"
putStrLn ""
putStrLn "Usage: periodic remove funcname jobname [jobname...]"
putStrLn ""
exitSuccess
printDropHelp :: IO ()
printDropHelp = do
putStrLn "periodic drop - Drop function"
putStrLn ""
putStrLn "Usage: periodic drop funcname [funcname...]"
putStrLn ""
exitSuccess
printConfigHelp :: IO ()
printConfigHelp = do
putStrLn "periodic config - Set or get config"
putStrLn ""
putStrLn "Usage: periodic config command [options]"
putStrLn ""
putStrLn "Commands:"
putStrLn " get Get config"
putStrLn " set Set config"
putStrLn ""
exitSuccess
printConfigKeyList :: IO ()
printConfigKeyList = do
putStrLn ""
putStrLn "Available keys:"
putStrLn " poll-interval - poll loop every time interval"
putStrLn " revert-interval - revert process queue loop every time interval"
putStrLn " timeout - job process timeout"
putStrLn " keepalive - client keepalive"
putStrLn " max-batch-size - max poll batch size"
putStrLn " expiration - run job result cache expiration"
putStrLn ""
exitSuccess
printConfigGetHelp :: IO ()
printConfigGetHelp = do
putStrLn "periodic config get - Get config"
putStrLn ""
putStrLn "Usage: periodic config get key"
printConfigKeyList
printConfigSetHelp :: IO ()
printConfigSetHelp = do
putStrLn "periodic config set - Set config"
putStrLn ""
putStrLn "Usage: periodic config set key val"
printConfigKeyList
printLoadHelp :: IO ()
printLoadHelp = do
putStrLn "periodic load - Load jobs"
putStrLn ""
putStrLn "Usage: periodic load file"
putStrLn ""
exitSuccess
printDumpHelp :: IO ()
printDumpHelp = do
putStrLn "periodic dump - Dump jobs"
putStrLn ""
putStrLn "Usage: periodic dump file"
putStrLn ""
exitSuccess
main :: IO ()
main = do
h <- lookupEnv "PERIODIC_PORT"
f <- lookupEnv "XOR_FILE"
(cmd, opts@Options {..}, argv) <- flip parseOptions (options h f) <$> getArgs
let argc = length argv
when (cmd == Help) printHelp
when (cmd == Submit && argc < 2) printSubmitHelp
when (cmd == Run && argc < 2) printRunHelp
when (cmd == Remove && argc < 2) printRemoveHelp
when (cmd == Drop && argc < 1) printDropHelp
when (cmd == Config && argc < 1) printConfigHelp
when (cmd == Load && argc < 1) printLoadHelp
when (cmd == Dump && argc < 1) printDumpHelp
when (not ("tcp" `isPrefixOf` host) && not ("unix" `isPrefixOf` host)) $ do
putStrLn $ "Invalid host " ++ host
printHelp
run opts cmd argv
run Options {useTls = True, ..} cmd argv = do
prms <- makeClientParams' cert [] certKey caStore (hostName, B.pack $ fromMaybe "" $ getService host)
clientEnv <- open (tlsConfig prms (socket host))
runClientT clientEnv $ processCommand cmd argv
run Options {useWs = True, ..} cmd argv = do
clientEnv <- open (clientConfig (socket host) (fromMaybe "0.0.0.0" $ getHost host) (fromMaybe "" $ getService host))
runClientT clientEnv $ processCommand cmd argv
run Options {xorFile = "", ..} cmd argv = do
clientEnv <- open (socket host)
runClientT clientEnv $ processCommand cmd argv
run Options {..} cmd argv = do
clientEnv <- open (xorConfig xorFile $ socket host)
runClientT clientEnv $ processCommand cmd argv
processCommand :: Transport tp => Command -> [String] -> ClientT tp IO ()
processCommand Help _ = liftIO printHelp
processCommand Status _ = doStatus
processCommand Submit xs = doSubmitJob xs
processCommand Run xs = doRunJob xs
processCommand Remove xs = doRemoveJob xs
processCommand Drop xs = doDropFunc xs
processCommand Config xs = doConfig xs
processCommand Load xs = doLoad xs
processCommand Dump xs = doDump xs
processCommand Shutdown _ = shutdown
doRemoveJob :: Transport tp => [String] -> ClientT tp IO ()
doRemoveJob (x:xs) = mapM_ (removeJob (fromString x) . fromString) xs
doRemoveJob [] = liftIO printRemoveHelp
doDropFunc :: Transport tp => [String] -> ClientT tp IO ()
doDropFunc = mapM_ (dropFunc . fromString)
doConfig :: Transport tp => [String] -> ClientT tp IO ()
doConfig ["get"] = liftIO printConfigGetHelp
doConfig ["get", k] = do
v <- configGet k
liftIO $ print v
doConfig ("get":_:_) = liftIO printConfigGetHelp
doConfig ["set"] = liftIO printConfigSetHelp
doConfig ["set", _] = liftIO printConfigSetHelp
doConfig ["set", k, v] = void $ configSet k (read v)
doConfig ("set":_:_:_) = liftIO printConfigSetHelp
doConfig _ = liftIO printConfigHelp
doLoad :: Transport tp => [String] -> ClientT tp IO ()
doLoad [fn] = do
jobs <- liftIO $ decodeFile fn
void $ load jobs
doLoad _ = liftIO printLoadHelp
doDump :: Transport tp => [String] -> ClientT tp IO ()
doDump [fn] = do
jobs <- dump
liftIO $ encodeFile fn jobs
doDump _ = liftIO printDumpHelp
doSubmitJob :: Transport tp => [String] -> ClientT tp IO ()
doSubmitJob [] = liftIO printSubmitHelp
doSubmitJob [_] = liftIO printSubmitHelp
doSubmitJob (x:y:xs) = do
(wl, l) <- liftIO $ go (Nothing, Nothing) xs
void $ submitJob (fromString x) (fromString y) wl l
where go :: (Maybe Workload, Maybe Int64) -> [String] -> IO (Maybe Workload, Maybe Int64)
go v [] = pure v
go (_, l0) ("-w":('@':f):ys) = do
w <- B.readFile f
go (Just (Workload w), l0) ys
go (_, l0) ("--workload":('@':f):ys) = do
w <- B.readFile f
go (Just (Workload w), l0) ys
go (_, l0) ("-w":w:ys) = go (Just (fromString w), l0) ys
go (_, l0) ("--workload":w:ys) = go (Just (fromString w), l0) ys
go (w, _) ("--later":l0:ys) = go (w, readMaybe l0) ys
go v (_:ys) = go v ys
doRunJob :: Transport tp => [String] -> ClientT tp IO ()
doRunJob [] = liftIO printRunHelp
doRunJob [_] = liftIO printRunHelp
doRunJob (x:y:xs) =
liftIO . putR
=<< runJob (fromString x) (fromString y)
=<< liftIO (go xs)
where go :: [String] -> IO (Maybe Workload)
go [] = pure Nothing
go ("-w":('@':f):_) = getFile f
go ("--workload":('@':f):_) = getFile f
go ("-w":w:_) = pure $ Just (fromString w)
go ("--workload":w:_) = pure $ Just (fromString w)
go (_:ys) = go ys
getFile :: String -> IO (Maybe Workload)
getFile fn = Just . Workload <$> B.readFile fn
putR :: Maybe ByteString -> IO ()
putR Nothing = putStrLn "Error: run job failed"
putR (Just bs) = B.putStr bs
doStatus :: Transport tp => ClientT tp IO ()
doStatus = do
st <- map formatTime . unpackBS . map (B.split ',') . B.lines <$> status
liftIO $ printTable (["FUNCTIONS", "WORKERS", "JOBS", "PROCESSING", "LOCKING", "SCHEDAT"]:st)
unpackBS :: [[ByteString]] -> [[String]]
unpackBS = map (map B.unpack)
formatTime :: [String] -> [String]
formatTime [] = []
formatTime [x] = [formatUnixTimeLocal x]
formatTime (x:xs) = x:formatTime xs
formatUnixTimeLocal :: String -> String
formatUnixTimeLocal = B.unpack
. unsafePerformIO
. formatUnixTime "%Y-%m-%d %H:%M:%S"
. fromEpochTime
. fromIntegral
. read
printTable :: [[String]] -> IO ()
printTable rows = T.printBox $ T.hsep 2 T.left (map (T.vcat T.left . map T.text) (transpose rows))