-- | This module contains the Haxl logic
module Haxible.DataSource (AnsibleHaxl, initHaxibleState, dataFetch, TaskReq (..)) where

import Control.Concurrent.Async (async)
import Control.Exception (Exception, SomeException, try)
import Control.Lens
import Data.Aeson.Lens
import Data.Hashable (Hashable (hashWithSalt))
import Data.List qualified as List
import Data.Text qualified as Text
import Data.Typeable (Typeable)
import Haxible.Connection (Connections (..), TaskCall (..))
import Haxible.Prelude hiding (State, state)
import Haxl.Core
import Say
import System.Clock qualified as Clock

type AnsibleHaxl a = GenHaxl () () a

data TaskReq a where
  RunTask :: TaskCall -> TaskReq Value
  deriving (Typeable)

deriving instance Eq (TaskReq a)

deriving instance Show (TaskReq a)

instance ShowP TaskReq where showp :: forall a. TaskReq a -> String
showp = TaskReq a -> String
forall a. Show a => a -> String
show

instance Hashable (TaskReq a) where
  hashWithSalt :: Int -> TaskReq a -> Int
hashWithSalt Int
s (RunTask TaskCall
param) = Int -> (Int, TaskCall) -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
s (Int
0 :: Int, TaskCall
param)

instance StateKey TaskReq where
  data State TaskReq = AnsibleState {State TaskReq -> Connections
connections :: Connections}

instance DataSourceName TaskReq where
  dataSourceName :: Proxy TaskReq -> Text
dataSourceName Proxy TaskReq
_ = Text
"Ansible"

instance DataSource u TaskReq where
  fetch :: State TaskReq -> Flags -> u -> PerformFetch TaskReq
fetch = State TaskReq -> Flags -> u -> PerformFetch TaskReq
forall u. State TaskReq -> Flags -> u -> PerformFetch TaskReq
fetchTask

initHaxibleState :: Connections -> IO (State TaskReq)
initHaxibleState :: Connections -> IO (State TaskReq)
initHaxibleState Connections
connections = State TaskReq -> IO (State TaskReq)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (State TaskReq -> IO (State TaskReq))
-> State TaskReq -> IO (State TaskReq)
forall a b. (a -> b) -> a -> b
$ AnsibleState {Connections
connections :: Connections
$sel:connections:AnsibleState :: Connections
connections}

fetchTask :: State TaskReq -> Flags -> u -> PerformFetch TaskReq
fetchTask :: forall u. State TaskReq -> Flags -> u -> PerformFetch TaskReq
fetchTask State TaskReq
state Flags
_flags u
_user =
  ([BlockedFetch TaskReq] -> IO ()) -> PerformFetch TaskReq
forall (req :: * -> *).
([BlockedFetch req] -> IO ()) -> PerformFetch req
BackgroundFetch (([BlockedFetch TaskReq] -> IO ()) -> PerformFetch TaskReq)
-> ([BlockedFetch TaskReq] -> IO ()) -> PerformFetch TaskReq
forall a b. (a -> b) -> a -> b
$ \[BlockedFetch TaskReq]
reqs -> do
    Maybe String
debug <- String -> IO (Maybe String)
lookupEnv String
"HAXIBLE_DEBUG"
    case Maybe String
debug of
      Just String
_ -> Text -> IO ()
forall (m :: * -> *). MonadIO m => Text -> m ()
say (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"▶ Batching " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
Text.pack (Int -> String
forall a. Show a => a -> String
show ([BlockedFetch TaskReq] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [BlockedFetch TaskReq]
reqs)) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" tasks"
      Maybe String
Nothing -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Integer
now <- TimeSpec -> Integer
Clock.toNanoSecs (TimeSpec -> Integer) -> IO TimeSpec -> IO Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Clock -> IO TimeSpec
Clock.getTime Clock
Clock.Monotonic
    (BlockedFetch TaskReq -> IO ()) -> [BlockedFetch TaskReq] -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Connections -> Integer -> BlockedFetch TaskReq -> IO ()
fetchAsync State TaskReq
state.connections Integer
now) [BlockedFetch TaskReq]
reqs

data TaskError = TaskError Int Value
  deriving (Int -> TaskError -> ShowS
[TaskError] -> ShowS
TaskError -> String
(Int -> TaskError -> ShowS)
-> (TaskError -> String)
-> ([TaskError] -> ShowS)
-> Show TaskError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TaskError] -> ShowS
$cshowList :: [TaskError] -> ShowS
show :: TaskError -> String
$cshow :: TaskError -> String
showsPrec :: Int -> TaskError -> ShowS
$cshowsPrec :: Int -> TaskError -> ShowS
Show)

instance Exception TaskError

fetchAsync :: Connections -> Integer -> BlockedFetch TaskReq -> IO ()
fetchAsync :: Connections -> Integer -> BlockedFetch TaskReq -> IO ()
fetchAsync Connections
python Integer
ts (BlockedFetch (RunTask TaskCall
task) ResultVar a
rvar) =
  IO (Async ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Async ()) -> IO ()) -> IO (Async ()) -> IO ()
forall a b. (a -> b) -> a -> b
$
    IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ do
      Either SomeException (Int, Value)
resultsE <- IO (Int, Value) -> IO (Either SomeException (Int, Value))
forall e a. Exception e => IO a -> IO (Either e a)
Control.Exception.try (IO (Int, Value) -> IO (Either SomeException (Int, Value)))
-> IO (Int, Value) -> IO (Either SomeException (Int, Value))
forall a b. (a -> b) -> a -> b
$ Connections
python.run TaskCall
task
      Integer
now <- TimeSpec -> Integer
Clock.toNanoSecs (TimeSpec -> Integer) -> IO TimeSpec -> IO Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Clock -> IO TimeSpec
Clock.getTime Clock
Clock.Monotonic
      case Either SomeException (Int, Value)
resultsE of
        Left SomeException
ex -> ResultVar a -> SomeException -> IO ()
forall e a. Exception e => ResultVar a -> e -> IO ()
putFailure ResultVar a
rvar (SomeException
ex :: SomeException)
        Right (Int
0, Value
result) -> ResultVar a -> a -> IO ()
forall a. ResultVar a -> a -> IO ()
putSuccess ResultVar a
rvar (Integer -> Value -> a
addTS Integer
now Value
result)
        Right (Int
code, Value
res) -> ResultVar a -> TaskError -> IO ()
forall e a. Exception e => ResultVar a -> e -> IO ()
putFailure ResultVar a
rvar (Int -> Value -> TaskError
TaskError Int
code (Integer -> Value -> a
addTS Integer
now Value
res))
  where
    addTS :: Integer -> Value -> a
addTS Integer
now =
      ((HashMap Text Value -> Identity (HashMap Text Value))
-> Value -> Identity a
forall t. AsValue t => Prism' t (HashMap Text Value)
_Object ((HashMap Text Value -> Identity (HashMap Text Value))
 -> Value -> Identity a)
-> ((Maybe (IxValue (HashMap Text Value))
     -> Identity (Maybe Value))
    -> HashMap Text Value -> Identity (HashMap Text Value))
-> (Maybe (IxValue (HashMap Text Value)) -> Identity (Maybe Value))
-> Value
-> Identity a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Index (HashMap Text Value)
-> Lens'
     (HashMap Text Value) (Maybe (IxValue (HashMap Text Value)))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at Text
Index (HashMap Text Value)
"__haxible_start" ((Maybe (IxValue (HashMap Text Value)) -> Identity (Maybe Value))
 -> Value -> Identity a)
-> Value -> Value -> a
forall s t a b. ASetter s t a (Maybe b) -> b -> s -> t
?~ Scientific -> Value
Number (Integer -> Scientific
forall a. Num a => Integer -> a
fromInteger Integer
ts))
        (Value -> a) -> (Value -> Value) -> Value -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((HashMap Text Value -> Identity (HashMap Text Value))
-> Value -> Identity Value
forall t. AsValue t => Prism' t (HashMap Text Value)
_Object ((HashMap Text Value -> Identity (HashMap Text Value))
 -> Value -> Identity Value)
-> ((Maybe (IxValue (HashMap Text Value))
     -> Identity (Maybe Value))
    -> HashMap Text Value -> Identity (HashMap Text Value))
-> (Maybe (IxValue (HashMap Text Value)) -> Identity (Maybe Value))
-> Value
-> Identity Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Index (HashMap Text Value)
-> Lens'
     (HashMap Text Value) (Maybe (IxValue (HashMap Text Value)))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at Text
Index (HashMap Text Value)
"__haxible_end" ((Maybe (IxValue (HashMap Text Value)) -> Identity (Maybe Value))
 -> Value -> Identity Value)
-> Value -> Value -> Value
forall s t a b. ASetter s t a (Maybe b) -> b -> s -> t
?~ Scientific -> Value
Number (Integer -> Scientific
forall a. Num a => Integer -> a
fromInteger Integer
now))