{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Main
( main
) where
import Control.Concurrent (forkIO, killThread)
import Control.DeepSeq (rnf)
import Control.Monad (unless, void, when)
import Control.Monad.IO.Class (liftIO)
import qualified Data.ByteString.Char8 as B (ByteString, pack)
import qualified Data.ByteString.Lazy as LB (null, toStrict)
import qualified Data.ByteString.Lazy.Char8 as LB (hGetContents, hPut)
import Data.List (isPrefixOf)
import Data.Maybe (fromMaybe)
import qualified Data.Text as T (unpack)
import Data.Text.Encoding (decodeUtf8With)
import Data.Text.Encoding.Error (ignore)
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.Job (JobT, name, withLock_, workDone,
workDone_, workFail, workload)
import Periodic.Trans.Worker (WorkerT, addFunc, broadcast,
startWorkerT, work)
import Periodic.Types (FuncName (..), LockName (..))
import System.Environment (getArgs, lookupEnv)
import System.Exit (ExitCode (..), exitSuccess)
import System.IO (hClose)
import System.Process (CreateProcess (std_in, std_out),
StdStream (CreatePipe, Inherit),
proc, waitForProcess,
withCreateProcess)
import UnliftIO (MVar, SomeException, evaluate,
mask, newEmptyMVar, onException,
putMVar, takeMVar, throwIO, try,
tryIO)
data Options = Options
{ host :: String
, xorFile :: FilePath
, useTls :: Bool
, useWs :: Bool
, hostName :: String
, certKey :: FilePath
, cert :: FilePath
, caStore :: FilePath
, thread :: Int
, lockCount :: Int
, lockName :: Maybe LockName
, notify :: Bool
, useData :: Bool
, useName :: Bool
, showHelp :: Bool
}
options :: Maybe Int -> Maybe String -> Maybe String -> Options
options t 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"
, thread = fromMaybe 1 t
, lockCount = 1
, lockName = Nothing
, notify = False
, useData = False
, useName = True
, showHelp = False
}
parseOptions :: [String] -> Options -> (Options, FuncName, String, [String])
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 ("--thread":x:xs) opt = parseOptions xs opt { thread = read x }
parseOptions ("--lock-count":x:xs) opt = parseOptions xs opt { lockCount = read x }
parseOptions ("--lock-name":x:xs) opt = parseOptions xs opt { lockName = Just (LockName $ B.pack x) }
parseOptions ("--help":xs) opt = parseOptions xs opt { showHelp = True }
parseOptions ("--broadcast":xs) opt = parseOptions xs opt { notify = True }
parseOptions ("--data":xs) opt = parseOptions xs opt { useData = True }
parseOptions ("--no-name":xs) opt = parseOptions xs opt { useName = False }
parseOptions ("-h":xs) opt = parseOptions xs opt { showHelp = True }
parseOptions [] opt = (opt { showHelp = True }, "", "", [])
parseOptions [_] opt = (opt { showHelp = True }, "", "", [])
parseOptions (x:y:xs) opt = (opt, FuncName $ B.pack x, y, xs)
printHelp :: IO ()
printHelp = do
putStrLn "periodic-run - Periodic task system worker"
putStrLn ""
putStrLn "Usage: periodic-run [--host|-H HOST] [--xor FILE|--ws|--tls [--hostname HOSTNAME] [--cert-key FILE] [--cert FILE] [--ca FILE] [--thread THREAD] [--lock-name NAME] [--lock-count COUNT] [--broadcast] [--data] [--no-name]] funcname command [options]"
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 " --thread Worker thread [$THREAD]"
putStrLn " --lock-count Max lock count (optional: 1)"
putStrLn " --lock-name The lock name (optional: no lock)"
putStrLn " --broadcast Is broadcast worker"
putStrLn " --data Send work data to client"
putStrLn " --no-name Ignore the job name"
putStrLn " -h --help Display help message"
putStrLn ""
putStrLn "Version: v1.1.7.1"
putStrLn ""
exitSuccess
main :: IO ()
main = do
h <- lookupEnv "PERIODIC_PORT"
f <- lookupEnv "XOR_FILE"
t <- fmap read <$> lookupEnv "THREAD"
(opts@Options {..}, func, cmd, argv) <- flip parseOptions (options t h f) <$> getArgs
when showHelp printHelp
when (not ("tcp" `isPrefixOf` host) && not ("unix" `isPrefixOf` host)) $ do
putStrLn $ "Invalid host " ++ host
printHelp
run opts func cmd argv
doWork :: Transport tp => Options -> FuncName -> String -> [String] -> WorkerT tp IO ()
doWork opts@Options{..} func cmd argv = do
let w = processWorker opts cmd argv
if notify then broadcast func w
else
case lockName of
Nothing -> addFunc func w
Just n -> addFunc func $ withLock_ n lockCount w
liftIO $ putStrLn "Worker started."
work thread
run :: Options -> FuncName -> String -> [String] -> IO ()
run opts@Options {useTls = True, ..} func cmd argv = do
prms <- makeClientParams' cert [] certKey caStore (hostName, B.pack $ fromMaybe "" $ getService host)
startWorkerT (tlsConfig prms (socket host)) $ doWork opts func cmd argv
run opts@Options {useWs = True, ..} func cmd argv =
startWorkerT (clientConfig (socket host) (fromMaybe "0.0.0.0" $ getHost host) (fromMaybe "" $ getService host)) $ doWork opts func cmd argv
run opts@Options {xorFile = "", ..} func cmd argv =
startWorkerT (socket host) $ doWork opts func cmd argv
run opts@Options {..} func cmd argv =
startWorkerT (xorConfig xorFile $ socket host) $ doWork opts func cmd argv
processWorker :: Transport tp => Options -> String -> [String] -> JobT tp IO ()
processWorker Options{..} cmd argv = do
n <- name
rb <- workload
let argv' = if useName then argv ++ [n] else argv
cp = (proc cmd argv') {std_in = CreatePipe, std_out= if useData then CreatePipe else Inherit}
(code, out) <- liftIO $ withCreateProcess cp $ \mb_inh mb_outh _ ph ->
case (mb_inh, mb_outh) of
(Nothing, _) -> error "processWorker: Failed to get a stdin handle."
(Just inh, Nothing) -> do
unless (LB.null rb) $ void $ tryIO $ LB.hPut inh rb
void $ tryIO $ hClose inh
code <- waitForProcess ph
return (code, Nothing)
(Just inh, Just outh) -> do
output <- LB.hGetContents outh
withForkWait (evaluate $ rnf output) $ \waitOut -> do
unless (LB.null rb) $ void $ tryIO $ LB.hPut inh rb
void $ tryIO $ hClose inh
waitOut
hClose outh
code <- waitForProcess ph
return (code, Just output)
case code of
ExitFailure _ -> workFail
ExitSuccess ->
case out of
Nothing -> workDone
Just wl -> workDone_ $ LB.toStrict wl
unpackBS :: B.ByteString -> String
unpackBS = T.unpack . decodeUtf8With ignore
-- | Fork a thread while doing something else, but kill it if there's an
-- exception.
--
-- This is important in the cases above because we want to kill the thread
-- that is holding the Handle lock, because when we clean up the process we
-- try to close that handle, which could otherwise deadlock.
--
withForkWait :: IO () -> (IO () -> IO a) -> IO a
withForkWait async body = do
waitVar <- newEmptyMVar :: IO (MVar (Either SomeException ()))
mask $ \restore -> do
tid <- forkIO $ try (restore async) >>= putMVar waitVar
let wait = takeMVar waitVar >>= either throwIO return
restore (body wait) `onException` killThread tid