module Haxible.Connection (Connections (..), TaskCall (..), withConnections, cleanVar) where
import Control.Exception (bracket)
import Data.Aeson (eitherDecodeStrict, encode)
import Data.Aeson.KeyMap qualified
import Data.ByteString (hGetLine, toStrict)
import Data.ByteString.Char8 (hPutStrLn)
import Data.Hashable (Hashable)
import Data.Pool qualified
import Data.Text qualified as Text
import Data.Text.Encoding (decodeUtf8)
import Data.Typeable (Typeable)
import GHC.IO.Handle (hFlush)
import Haxible.Prelude
import Say
import System.Console.ANSI (Color (..), ColorIntensity (Dull), ConsoleLayer (Foreground), hSupportsANSIColor, setSGRCode)
import System.Console.ANSI.Codes (SGR (..))
import System.IO (Handle, hClose, stdout)
import System.Process (Pid, getPid)
import System.Process.Typed
data TaskCall = TaskCall
{
TaskCall -> Vars
playAttrs :: Vars,
TaskCall -> Text
module_ :: Text,
TaskCall -> Value
moduleObject :: Value,
TaskCall -> Vars
taskAttrs :: Vars,
TaskCall -> Vars
taskVars :: Vars
}
deriving (TaskCall -> TaskCall -> Bool
(TaskCall -> TaskCall -> Bool)
-> (TaskCall -> TaskCall -> Bool) -> Eq TaskCall
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: TaskCall -> TaskCall -> Bool
$c/= :: TaskCall -> TaskCall -> Bool
== :: TaskCall -> TaskCall -> Bool
$c== :: TaskCall -> TaskCall -> Bool
Eq, Int -> TaskCall -> ShowS
[TaskCall] -> ShowS
TaskCall -> String
(Int -> TaskCall -> ShowS)
-> (TaskCall -> String) -> ([TaskCall] -> ShowS) -> Show TaskCall
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TaskCall] -> ShowS
$cshowList :: [TaskCall] -> ShowS
show :: TaskCall -> String
$cshow :: TaskCall -> String
showsPrec :: Int -> TaskCall -> ShowS
$cshowsPrec :: Int -> TaskCall -> ShowS
Show, Typeable, (forall x. TaskCall -> Rep TaskCall x)
-> (forall x. Rep TaskCall x -> TaskCall) -> Generic TaskCall
forall x. Rep TaskCall x -> TaskCall
forall x. TaskCall -> Rep TaskCall x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep TaskCall x -> TaskCall
$cfrom :: forall x. TaskCall -> Rep TaskCall x
Generic, Eq TaskCall
Eq TaskCall
-> (Int -> TaskCall -> Int)
-> (TaskCall -> Int)
-> Hashable TaskCall
Int -> TaskCall -> Int
TaskCall -> Int
forall a. Eq a -> (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: TaskCall -> Int
$chash :: TaskCall -> Int
hashWithSalt :: Int -> TaskCall -> Int
$chashWithSalt :: Int -> TaskCall -> Int
Hashable)
newtype Connections = Connections {Connections -> TaskCall -> IO (Int, Value)
run :: TaskCall -> IO (Int, Value)}
addSep :: Int -> Text -> Text
addSep :: Int -> Text -> Text
addSep Int
width Text
x = Text
x Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
sep
where
sep :: Text
sep = Int -> Text -> Text
Text.replicate (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- Text -> Int
Text.length Text
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Text
"*"
formatPid :: Pid -> Text
formatPid :: Pid -> Text
formatPid Pid
pid = Text
"<" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
forall source target. From source target => source -> target
from (Pid -> String
forall a. Show a => a -> String
show Pid
pid) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
">"
formatTask :: Pid -> TaskCall -> Text
formatTask :: Pid -> TaskCall -> Text
formatTask Pid
pid TaskCall
tc = Text
"TASK [" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"] " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Pid -> Text
formatPid Pid
pid
where
name :: Text
name = Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe TaskCall
tc.module_ (Getting (First Text) Value Text -> Value -> Maybe Text
forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview (Text -> Traversal' Value Value
forall t. AsValue t => Text -> Traversal' t Value
key Text
"name" ((Value -> Const (First Text) Value)
-> Value -> Const (First Text) Value)
-> Getting (First Text) Value Text
-> Getting (First Text) Value Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (First Text) Value Text
forall t. AsPrimitive t => Prism' t Text
_String) TaskCall
tc.moduleObject)
formatResult :: Bool -> Pid -> (Int, Value) -> String
formatResult :: Bool -> Pid -> (Int, Value) -> String
formatResult Bool
withColor Pid
pid (Int
code, Value
val) = String
pre String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Text -> String
forall source target. From source target => source -> target
from Text
txt String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
post
where
(String
pre, String
post)
| Bool
withColor = ([SGR] -> String
setSGRCode [ConsoleLayer -> ColorIntensity -> Color -> SGR
SetColor ConsoleLayer
Foreground ColorIntensity
Dull Color
resColor], [SGR] -> String
setSGRCode [SGR
Reset])
| Bool
otherwise = (String
"", String
"")
txt :: Text
txt = Text
res Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": [" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
host Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"] " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Pid -> Text
formatPid Pid
pid Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" => " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
jsonDump Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n"
jsonDump :: Text
jsonDump = ByteString -> Text
decodeUtf8 (ByteString -> Text) -> (Value -> ByteString) -> Value -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
forall source target. From source target => source -> target
from (ByteString -> ByteString)
-> (Value -> ByteString) -> Value -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> ByteString
forall a. ToJSON a => a -> ByteString
encode (Value -> ByteString) -> (Value -> Value) -> Value -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Value
cleanVar (Value -> Text) -> Value -> Text
forall a b. (a -> b) -> a -> b
$ Value
val
host :: Text
host = Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
"unknown?!" (Getting (First Text) Value Text -> Value -> Maybe Text
forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview (Text -> Traversal' Value Value
forall t. AsValue t => Text -> Traversal' t Value
key Text
"__haxible_play" ((Value -> Const (First Text) Value)
-> Value -> Const (First Text) Value)
-> Getting (First Text) Value Text
-> Getting (First Text) Value Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Traversal' Value Value
forall t. AsValue t => Text -> Traversal' t Value
key Text
"hosts" ((Value -> Const (First Text) Value)
-> Value -> Const (First Text) Value)
-> Getting (First Text) Value Text
-> Getting (First Text) Value Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (First Text) Value Text
forall t. AsPrimitive t => Prism' t Text
_String) Value
val)
resColor :: Color
resColor
| Int
code Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 = Color
Red
| Bool
otherwise = Color
Green
res :: Text
res
| Int
code Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Text
"ok"
| Bool
otherwise = String -> Text
forall source target. From source target => source -> target
from (Int -> String
forall a. Show a => a -> String
show Int
code)
cleanVar :: Value -> Value
cleanVar :: Value -> Value
cleanVar = \case
Object Object
obj -> Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ (Key -> Value -> Bool) -> Object -> Object
forall v. (Key -> v -> Bool) -> KeyMap v -> KeyMap v
Data.Aeson.KeyMap.filterWithKey (\Key
k Value
_ -> Key
k Key -> [Key] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Key]
addedKey) Object
obj
Value
x -> Value
x
where
addedKey :: [Key]
addedKey =
[Key
"__haxible_play", Key
"__haxible_start", Key
"__haxible_end", Key
"__haxible_module"]
[Key] -> [Key] -> [Key]
forall a. Semigroup a => a -> a -> a
<> [Key
"__haxible_multi_hosts"]
[Key] -> [Key] -> [Key]
forall a. Semigroup a => a -> a -> a
<> [Key
"_ansible_no_log", Key
"_ansible_verbose_always"]
withConnections :: Int -> FilePath -> (Connections -> IO ()) -> IO ()
withConnections :: Int -> String -> (Connections -> IO ()) -> IO ()
withConnections Int
count String
inventory Connections -> IO ()
callback =
IO (Pool (Process Handle Handle ()))
-> (Pool (Process Handle Handle ()) -> IO ())
-> (Pool (Process Handle Handle ()) -> IO ())
-> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (PoolConfig (Process Handle Handle ())
-> IO (Pool (Process Handle Handle ()))
forall a. PoolConfig a -> IO (Pool a)
Data.Pool.newPool PoolConfig (Process Handle Handle ())
poolConfig) Pool (Process Handle Handle ()) -> IO ()
forall a. Pool a -> IO ()
Data.Pool.destroyAllResources Pool (Process Handle Handle ()) -> IO ()
go
where
go :: Pool (Process Handle Handle ()) -> IO ()
go Pool (Process Handle Handle ())
pool = do
Int
termWidth <- Int -> (String -> Int) -> Maybe String -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
80 String -> Int
forall a. Read a => String -> a
read (Maybe String -> Int) -> IO (Maybe String) -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Maybe String)
lookupEnv String
"COLUMNS"
Bool
withColor <- Handle -> IO Bool
hSupportsANSIColor Handle
stdout
Text -> IO ()
forall (m :: * -> *). MonadIO m => Text -> m ()
say (Int -> Text -> Text
addSep Int
termWidth Text
"PLAY [concurrent]" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n")
let runTask :: TaskCall -> Process Handle Handle () -> IO (Int, Value)
runTask :: TaskCall -> Process Handle Handle () -> IO (Int, Value)
runTask TaskCall
taskCall Process Handle Handle ()
p = do
let callParams :: [Value]
callParams = [Vars -> Value
mkObj TaskCall
taskCall.playAttrs, TaskCall
taskCall.moduleObject, Vars -> Value
mkObj TaskCall
taskCall.taskAttrs, Vars -> Value
mkObj TaskCall
taskCall.taskVars]
Pid
pid <- Pid -> Maybe Pid -> Pid
forall a. a -> Maybe a -> a
fromMaybe (String -> Pid
forall a. HasCallStack => String -> a
error String
"no pid?!") (Maybe Pid -> Pid) -> IO (Maybe Pid) -> IO Pid
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ProcessHandle -> IO (Maybe Pid)
getPid (Process Handle Handle () -> ProcessHandle
forall stdin stdout stderr.
Process stdin stdout stderr -> ProcessHandle
unsafeProcessHandle Process Handle Handle ()
p)
Text -> IO ()
forall (m :: * -> *). MonadIO m => Text -> m ()
say (Int -> Text -> Text
addSep Int
termWidth (Pid -> TaskCall -> Text
formatTask Pid
pid TaskCall
taskCall))
Handle -> ByteString -> IO ()
hPutStrLn (Process Handle Handle () -> Handle
forall stdin stdout stderr. Process stdin stdout stderr -> stdin
getStdin Process Handle Handle ()
p) (ByteString -> ByteString
toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ [Value] -> ByteString
forall a. ToJSON a => a -> ByteString
encode [Value]
callParams)
Handle -> IO ()
hFlush (Process Handle Handle () -> Handle
forall stdin stdout stderr. Process stdin stdout stderr -> stdin
getStdin Process Handle Handle ()
p)
ByteString
output <- Handle -> IO ByteString
hGetLine (Process Handle Handle () -> Handle
forall stdin stdout stderr. Process stdin stdout stderr -> stdout
getStdout Process Handle Handle ()
p)
case ByteString -> Either String (Int, Value)
forall a. FromJSON a => ByteString -> Either String a
eitherDecodeStrict ByteString
output of
Right (Int, Value)
res -> do
String -> IO ()
forall (m :: * -> *). MonadIO m => String -> m ()
sayString (Bool -> Pid -> (Int, Value) -> String
formatResult Bool
withColor Pid
pid (Int, Value)
res)
(Int, Value) -> IO (Int, Value)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int, Value)
res
Left String
err -> String -> IO (Int, Value)
forall a. HasCallStack => String -> a
error (String -> IO (Int, Value)) -> String -> IO (Int, Value)
forall a b. (a -> b) -> a -> b
$ ByteString -> String
forall a. Show a => a -> String
show ByteString
output String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
": " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
err
Connections -> IO ()
callback ((TaskCall -> IO (Int, Value)) -> Connections
Connections ((TaskCall -> IO (Int, Value)) -> Connections)
-> (TaskCall -> IO (Int, Value)) -> Connections
forall a b. (a -> b) -> a -> b
$ Pool (Process Handle Handle ())
-> (Process Handle Handle () -> IO (Int, Value)) -> IO (Int, Value)
forall a r. Pool a -> (a -> IO r) -> IO r
Data.Pool.withResource Pool (Process Handle Handle ())
pool ((Process Handle Handle () -> IO (Int, Value)) -> IO (Int, Value))
-> (TaskCall -> Process Handle Handle () -> IO (Int, Value))
-> TaskCall
-> IO (Int, Value)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TaskCall -> Process Handle Handle () -> IO (Int, Value)
runTask)
poolConfig :: Data.Pool.PoolConfig (Process Handle Handle ())
poolConfig :: PoolConfig (Process Handle Handle ())
poolConfig =
Data.Pool.PoolConfig
{ createResource :: IO (Process Handle Handle ())
createResource = do
ProcessConfig Handle Handle () -> IO (Process Handle Handle ())
forall (m :: * -> *) stdin stdout stderr.
MonadIO m =>
ProcessConfig stdin stdout stderr
-> m (Process stdin stdout stderr)
startProcess
(ProcessConfig Handle Handle () -> IO (Process Handle Handle ()))
-> (ProcessConfig () () () -> ProcessConfig Handle Handle ())
-> ProcessConfig () () ()
-> IO (Process Handle Handle ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StreamSpec 'STInput Handle
-> ProcessConfig () Handle () -> ProcessConfig Handle Handle ()
forall stdin stdin0 stdout stderr.
StreamSpec 'STInput stdin
-> ProcessConfig stdin0 stdout stderr
-> ProcessConfig stdin stdout stderr
setStdin StreamSpec 'STInput Handle
forall (anyStreamType :: StreamType).
StreamSpec anyStreamType Handle
createPipe
(ProcessConfig () Handle () -> ProcessConfig Handle Handle ())
-> (ProcessConfig () () () -> ProcessConfig () Handle ())
-> ProcessConfig () () ()
-> ProcessConfig Handle Handle ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StreamSpec 'STOutput Handle
-> ProcessConfig () () () -> ProcessConfig () Handle ()
forall stdout stdin stdout0 stderr.
StreamSpec 'STOutput stdout
-> ProcessConfig stdin stdout0 stderr
-> ProcessConfig stdin stdout stderr
setStdout StreamSpec 'STOutput Handle
forall (anyStreamType :: StreamType).
StreamSpec anyStreamType Handle
createPipe
(ProcessConfig () () () -> IO (Process Handle Handle ()))
-> ProcessConfig () () () -> IO (Process Handle Handle ())
forall a b. (a -> b) -> a -> b
$ String -> [String] -> ProcessConfig () () ()
proc String
"python" [String
"./app/wrapper.py", String
inventory],
freeResource :: Process Handle Handle () -> IO ()
freeResource = \Process Handle Handle ()
p -> do
Handle -> IO ()
hClose (Process Handle Handle () -> Handle
forall stdin stdout stderr. Process stdin stdout stderr -> stdin
getStdin Process Handle Handle ()
p)
Process Handle Handle () -> IO ()
forall (m :: * -> *) stdin stdout stderr.
MonadIO m =>
Process stdin stdout stderr -> m ()
stopProcess Process Handle Handle ()
p,
poolCacheTTL :: Double
poolCacheTTL = Double
3600,
poolMaxResources :: Int
poolMaxResources = Int
count
}