module Hercules.Agent.WorkerProcess
  ( runWorker,
    getWorkerExe,
    getDaemonExe,
    WorkerEnvSettings (..),
    prepareEnv,
    modifyEnv,
  )
where

import Conduit hiding (Producer)
import qualified Control.Exception.Safe as Safe
import Control.Monad.IO.Unlift
import Data.Binary (Binary)
import Data.Conduit.Extras (conduitToCallbacks, sourceChan)
import Data.Conduit.Serialization.Binary
  ( conduitDecode,
    conduitEncode,
  )
import qualified Data.Map as M
import GHC.IO.Exception
import qualified Hercules.API.Agent.Evaluate.EvaluateTask as EvaluateTask
import Hercules.Agent.NixPath (renderNixPath)
import Paths_hercules_ci_agent (getBinDir)
import Protolude
import System.Environment (getEnvironment)
import System.FilePath ((</>))
import System.IO (hClose)
import System.IO.Error
import System.Process
import System.Timeout (timeout)
import Prelude ()

data WorkerException = WorkerException
  { WorkerException -> SomeException
originalException :: SomeException,
    WorkerException -> Maybe ExitCode
exitStatus :: Maybe ExitCode
  }
  deriving (Int -> WorkerException -> ShowS
[WorkerException] -> ShowS
WorkerException -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [WorkerException] -> ShowS
$cshowList :: [WorkerException] -> ShowS
show :: WorkerException -> [Char]
$cshow :: WorkerException -> [Char]
showsPrec :: Int -> WorkerException -> ShowS
$cshowsPrec :: Int -> WorkerException -> ShowS
Show, Typeable)

instance Exception WorkerException where
  displayException :: WorkerException -> [Char]
displayException WorkerException
we =
    forall e. Exception e => e -> [Char]
displayException (WorkerException -> SomeException
originalException WorkerException
we)
      forall a. Semigroup a => a -> a -> a
<> case WorkerException -> Maybe ExitCode
exitStatus WorkerException
we of
        Maybe ExitCode
Nothing -> [Char]
""
        Just ExitCode
s -> [Char]
" (worker: " forall a. Semigroup a => a -> a -> a
<> forall a b. (Show a, StringConv [Char] b) => a -> b
show ExitCode
s forall a. Semigroup a => a -> a -> a
<> [Char]
")"

data WorkerEnvSettings = WorkerEnvSettings
  { WorkerEnvSettings -> [NixPathElement (SubPathOf [Char])]
nixPath :: [EvaluateTask.NixPathElement (EvaluateTask.SubPathOf FilePath)],
    WorkerEnvSettings -> Map [Char] [Char]
extraEnv :: Map [Char] [Char]
  }

-- | Filter out impure env vars by wildcard, set NIX_PATH
modifyEnv :: WorkerEnvSettings -> Map [Char] [Char] -> Map [Char] [Char]
modifyEnv :: WorkerEnvSettings -> Map [Char] [Char] -> Map [Char] [Char]
modifyEnv WorkerEnvSettings
workerEnvSettings =
  forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union (WorkerEnvSettings -> Map [Char] [Char]
extraEnv WorkerEnvSettings
workerEnvSettings)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert [Char]
"NIX_PATH" (forall a b. ConvertText a b => a -> b
toS forall a b. (a -> b) -> a -> b
$ [NixPathElement (SubPathOf [Char])] -> Text
renderNixPath forall a b. (a -> b) -> a -> b
$ WorkerEnvSettings -> [NixPathElement (SubPathOf [Char])]
nixPath WorkerEnvSettings
workerEnvSettings)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\[Char]
k [Char]
_v -> Bool -> Bool
not ([Char]
"NIXPKGS_" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Char]
k))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\[Char]
k [Char]
_v -> Bool -> Bool
not ([Char]
"NIXOS_" forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Char]
k))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => k -> Map k a -> Map k a
M.delete [Char]
"IN_NIX_SHELL"

prepareEnv :: WorkerEnvSettings -> IO [([Char], [Char])]
prepareEnv :: WorkerEnvSettings -> IO [([Char], [Char])]
prepareEnv WorkerEnvSettings
workerEnvSettings = do
  Map [Char] [Char]
envMap <- forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO [([Char], [Char])]
getEnvironment
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$ WorkerEnvSettings -> Map [Char] [Char] -> Map [Char] [Char]
modifyEnv WorkerEnvSettings
workerEnvSettings Map [Char] [Char]
envMap

getWorkerExe :: MonadIO m => m [Char]
getWorkerExe :: forall (m :: * -> *). MonadIO m => m [Char]
getWorkerExe = do
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO [Char]
getBinDir forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> ([Char] -> ShowS
</> [Char]
"hercules-ci-agent-worker")

getDaemonExe :: MonadIO m => m [Char]
getDaemonExe :: forall (m :: * -> *). MonadIO m => m [Char]
getDaemonExe = do
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO [Char]
getBinDir forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> ([Char] -> ShowS
</> [Char]
"hercules-ci-nix-daemon")

-- | Control a child process by communicating over stdin and stdout
-- using a 'Binary' interface.
runWorker ::
  (Binary command, Binary event, MonadUnliftIO m, MonadThrow m) =>
  -- | Process invocation details. Will ignore std_in, std_out and std_err fields.
  CreateProcess ->
  (Int -> ByteString -> m ()) ->
  Chan (Maybe command) ->
  (event -> m ()) ->
  m ExitCode
runWorker :: forall command event (m :: * -> *).
(Binary command, Binary event, MonadUnliftIO m, MonadThrow m) =>
CreateProcess
-> (Int -> ByteString -> m ())
-> Chan (Maybe command)
-> (event -> m ())
-> m ExitCode
runWorker CreateProcess
baseProcess Int -> ByteString -> m ()
stderrLineHandler Chan (Maybe command)
commandChan event -> m ()
eventHandler = do
  UnliftIO {unliftIO :: forall (m :: * -> *). UnliftIO m -> forall a. m a -> IO a
unliftIO = forall a. m a -> IO a
unlift} <- forall (m :: * -> *). MonadUnliftIO m => m (UnliftIO m)
askUnliftIO
  let createProcessSpec :: CreateProcess
createProcessSpec =
        CreateProcess
baseProcess
          { std_in :: StdStream
std_in = StdStream
CreatePipe,
            std_out :: StdStream
std_out = StdStream
CreatePipe,
            std_err :: StdStream
std_err = StdStream
CreatePipe
          }
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$
    forall a.
CreateProcess
-> (Maybe Handle
    -> Maybe Handle -> Maybe Handle -> ProcessHandle -> IO a)
-> IO a
withCreateProcess CreateProcess
createProcessSpec forall a b. (a -> b) -> a -> b
$ \Maybe Handle
mIn Maybe Handle
mOut Maybe Handle
mErr ProcessHandle
processHandle -> do
      (Handle
inHandle, Handle
outHandle, Handle
errHandle) <-
        case (,,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Handle
mIn forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Handle
mOut forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Handle
mErr of
          Just (Handle, Handle, Handle)
x -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Handle, Handle, Handle)
x
          Maybe (Handle, Handle, Handle)
Nothing ->
            forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO forall a b. (a -> b) -> a -> b
$
              IOErrorType -> [Char] -> Maybe Handle -> Maybe [Char] -> IOError
mkIOError
                IOErrorType
illegalOperationErrorType
                [Char]
"Process did not return all handles"
                forall a. Maybe a
Nothing -- no handle
                forall a. Maybe a
Nothing -- no path
      Maybe Pid
pidMaybe <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ ProcessHandle -> IO (Maybe Pid)
getPid ProcessHandle
processHandle
      let pid :: Int
pid = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 forall a b. (Integral a, Num b) => a -> b
fromIntegral Maybe Pid
pidMaybe
      let stderrPiper :: IO ()
stderrPiper =
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$
              forall (m :: * -> *) r. Monad m => ConduitT () Void m r -> m r
runConduit
                (forall (m :: * -> *) i.
MonadIO m =>
Handle -> ConduitT i ByteString m ()
sourceHandle Handle
errHandle forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall (m :: * -> *) seq.
(Monad m, IsSequence seq, Element seq ~ Word8) =>
ConduitT seq seq m ()
linesUnboundedAsciiC forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall (m :: * -> *) i o r.
Monad m =>
(i -> ConduitT i o m r) -> ConduitT i o m ()
awaitForever (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. m a -> IO a
unlift forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> m ()
stderrLineHandler Int
pid))
      let eventConduit :: ConduitT a event m ()
eventConduit = forall (m :: * -> *) i.
MonadIO m =>
Handle -> ConduitT i ByteString m ()
sourceHandle Handle
outHandle forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall b (m :: * -> *).
(Binary b, MonadThrow m) =>
ConduitT ByteString b m ()
conduitDecode
          commandConduit :: ConduitT a c IO ()
commandConduit =
            forall (m :: * -> *) a i.
MonadIO m =>
Chan (Maybe a) -> ConduitT i a m ()
sourceChan Chan (Maybe command)
commandChan
              forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall b (m :: * -> *).
(Binary b, MonadThrow m) =>
ConduitT b ByteString m ()
conduitEncode
              forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall (m :: * -> *) mono a.
(Monad m, MonoFoldable mono) =>
(a -> mono) -> ConduitT a (Element mono) m ()
concatMapC (\ByteString
x -> [forall a. a -> Flush a
Chunk ByteString
x, forall a. Flush a
Flush])
              forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall (m :: * -> *) e i o r.
(MonadUnliftIO m, Exception e) =>
(e -> ConduitT i o m r) -> ConduitT i o m r -> ConduitT i o m r
handleC forall {f :: * -> *}. MonadIO f => IOError -> f ()
handleEPIPE (forall (m :: * -> *) o.
MonadIO m =>
Handle -> ConduitM (Flush ByteString) o m ()
sinkHandleFlush Handle
inHandle)
          handleEPIPE :: IOError -> f ()
handleEPIPE IOError
e | IOError -> IOErrorType
ioeGetErrorType IOError
e forall a. Eq a => a -> a -> Bool
== IOErrorType
ResourceVanished = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          handleEPIPE IOError
e = forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO IOError
e
      let cmdThread :: IO ()
cmdThread = forall (m :: * -> *) r. Monad m => ConduitT () Void m r -> m r
runConduit forall {a} {c}. ConduitT a c IO ()
commandConduit forall a b. IO a -> IO b -> IO a
`finally` Handle -> IO ()
hClose Handle
outHandle
          eventThread :: IO ()
eventThread = forall a. m a -> IO a
unlift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) o a.
(MonadUnliftIO m, MonadIO m) =>
ConduitT () o m a -> (o -> m ()) -> m a
conduitToCallbacks forall {a}. ConduitT a event m ()
eventConduit event -> m ()
eventHandler
      -- plain forkIO so it can process all of stderr in case of an exception
      forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO IO ()
stderrPiper
      forall a b. IO a -> (Async a -> IO b) -> IO b
withAsync (ProcessHandle -> IO ExitCode
waitForProcess ProcessHandle
processHandle) forall a b. (a -> b) -> a -> b
$ \Async ExitCode
exitAsync -> do
        forall a b. IO a -> (Async a -> IO b) -> IO b
withAsync IO ()
cmdThread forall a b. (a -> b) -> a -> b
$ \Async ()
_ -> do
          IO ()
eventThread
            forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> (e -> m a) -> m a
`Safe.catch` \SomeException
e -> do
              let oneSecond :: Int
oneSecond = Int
1000 forall a. Num a => a -> a -> a
* Int
1000
              Maybe ExitCode
maybeStatus <- forall a. Int -> IO a -> IO (Maybe a)
timeout (Int
5 forall a. Num a => a -> a -> a
* Int
oneSecond) (forall a. Async a -> IO a
wait Async ExitCode
exitAsync)
              forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO forall a b. (a -> b) -> a -> b
$ SomeException -> Maybe ExitCode -> WorkerException
WorkerException SomeException
e Maybe ExitCode
maybeStatus
          forall a. Async a -> IO a
wait Async ExitCode
exitAsync