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]
(Int -> WorkerException -> ShowS)
-> (WorkerException -> [Char])
-> ([WorkerException] -> ShowS)
-> Show WorkerException
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 =
    SomeException -> [Char]
forall e. Exception e => e -> [Char]
displayException (WorkerException -> SomeException
originalException WorkerException
we)
      [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> case WorkerException -> Maybe ExitCode
exitStatus WorkerException
we of
        Maybe ExitCode
Nothing -> [Char]
""
        Just ExitCode
s -> [Char]
" (worker: " [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> ExitCode -> [Char]
forall a b. (Show a, StringConv [Char] b) => a -> b
show ExitCode
s [Char] -> ShowS
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 =
  Map [Char] [Char] -> Map [Char] [Char] -> Map [Char] [Char]
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union (WorkerEnvSettings -> Map [Char] [Char]
extraEnv WorkerEnvSettings
workerEnvSettings)
    (Map [Char] [Char] -> Map [Char] [Char])
-> (Map [Char] [Char] -> Map [Char] [Char])
-> Map [Char] [Char]
-> Map [Char] [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> [Char] -> Map [Char] [Char] -> Map [Char] [Char]
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert [Char]
"NIX_PATH" (Text -> [Char]
forall a b. ConvertText a b => a -> b
toS (Text -> [Char]) -> Text -> [Char]
forall a b. (a -> b) -> a -> b
$ [NixPathElement (SubPathOf [Char])] -> Text
renderNixPath ([NixPathElement (SubPathOf [Char])] -> Text)
-> [NixPathElement (SubPathOf [Char])] -> Text
forall a b. (a -> b) -> a -> b
$ WorkerEnvSettings -> [NixPathElement (SubPathOf [Char])]
nixPath WorkerEnvSettings
workerEnvSettings)
    (Map [Char] [Char] -> Map [Char] [Char])
-> (Map [Char] [Char] -> Map [Char] [Char])
-> Map [Char] [Char]
-> Map [Char] [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Char] -> [Char] -> Bool)
-> Map [Char] [Char] -> Map [Char] [Char]
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\[Char]
k [Char]
_v -> Bool -> Bool
not ([Char]
"NIXPKGS_" [Char] -> [Char] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Char]
k))
    (Map [Char] [Char] -> Map [Char] [Char])
-> (Map [Char] [Char] -> Map [Char] [Char])
-> Map [Char] [Char]
-> Map [Char] [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Char] -> [Char] -> Bool)
-> Map [Char] [Char] -> Map [Char] [Char]
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\[Char]
k [Char]
_v -> Bool -> Bool
not ([Char]
"NIXOS_" [Char] -> [Char] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [Char]
k))
    (Map [Char] [Char] -> Map [Char] [Char])
-> (Map [Char] [Char] -> Map [Char] [Char])
-> Map [Char] [Char]
-> Map [Char] [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Map [Char] [Char] -> Map [Char] [Char]
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 <- [([Char], [Char])] -> Map [Char] [Char]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([([Char], [Char])] -> Map [Char] [Char])
-> IO [([Char], [Char])] -> IO (Map [Char] [Char])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO [([Char], [Char])]
getEnvironment
  [([Char], [Char])] -> IO [([Char], [Char])]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([([Char], [Char])] -> IO [([Char], [Char])])
-> [([Char], [Char])] -> IO [([Char], [Char])]
forall a b. (a -> b) -> a -> b
$ Map [Char] [Char] -> [([Char], [Char])]
forall k a. Map k a -> [(k, a)]
M.toList (Map [Char] [Char] -> [([Char], [Char])])
-> Map [Char] [Char] -> [([Char], [Char])]
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
  IO [Char] -> m [Char]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO [Char]
getBinDir m [Char] -> ShowS -> m [Char]
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
  IO [Char] -> m [Char]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO [Char]
getBinDir m [Char] -> ShowS -> m [Char]
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} <- m (UnliftIO m)
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
          }
  IO ExitCode -> m ExitCode
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ExitCode -> m ExitCode) -> IO ExitCode -> m ExitCode
forall a b. (a -> b) -> a -> b
$
    CreateProcess
-> (Maybe Handle
    -> Maybe Handle -> Maybe Handle -> ProcessHandle -> IO ExitCode)
-> IO ExitCode
forall a.
CreateProcess
-> (Maybe Handle
    -> Maybe Handle -> Maybe Handle -> ProcessHandle -> IO a)
-> IO a
withCreateProcess CreateProcess
createProcessSpec ((Maybe Handle
  -> Maybe Handle -> Maybe Handle -> ProcessHandle -> IO ExitCode)
 -> IO ExitCode)
-> (Maybe Handle
    -> Maybe Handle -> Maybe Handle -> ProcessHandle -> IO ExitCode)
-> IO ExitCode
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 (,,) (Handle -> Handle -> Handle -> (Handle, Handle, Handle))
-> Maybe Handle
-> Maybe (Handle -> Handle -> (Handle, Handle, Handle))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Handle
mIn Maybe (Handle -> Handle -> (Handle, Handle, Handle))
-> Maybe Handle -> Maybe (Handle -> (Handle, Handle, Handle))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Handle
mOut Maybe (Handle -> (Handle, Handle, Handle))
-> Maybe Handle -> Maybe (Handle, Handle, Handle)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Handle
mErr of
          Just (Handle, Handle, Handle)
x -> (Handle, Handle, Handle) -> IO (Handle, Handle, Handle)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Handle, Handle, Handle)
x
          Maybe (Handle, Handle, Handle)
Nothing ->
            IOError -> IO (Handle, Handle, Handle)
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO (IOError -> IO (Handle, Handle, Handle))
-> IOError -> IO (Handle, Handle, Handle)
forall a b. (a -> b) -> a -> b
$
              IOErrorType -> [Char] -> Maybe Handle -> Maybe [Char] -> IOError
mkIOError
                IOErrorType
illegalOperationErrorType
                [Char]
"Process did not return all handles"
                Maybe Handle
forall a. Maybe a
Nothing -- no handle
                Maybe [Char]
forall a. Maybe a
Nothing -- no path
      Maybe Pid
pidMaybe <- IO (Maybe Pid) -> IO (Maybe Pid)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe Pid) -> IO (Maybe Pid))
-> IO (Maybe Pid) -> IO (Maybe Pid)
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> IO (Maybe Pid)
getPid ProcessHandle
processHandle
      let pid :: Int
pid = Int -> (Pid -> Int) -> Maybe Pid -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 Pid -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Maybe Pid
pidMaybe
      let stderrPiper :: IO ()
stderrPiper =
            IO () -> IO ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
              ConduitT () Void IO () -> IO ()
forall (m :: * -> *) r. Monad m => ConduitT () Void m r -> m r
runConduit
                (Handle -> ConduitT () ByteString IO ()
forall (m :: * -> *) i.
MonadIO m =>
Handle -> ConduitT i ByteString m ()
sourceHandle Handle
errHandle ConduitT () ByteString IO ()
-> ConduitM ByteString Void IO () -> ConduitT () Void IO ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitT ByteString ByteString IO ()
forall (m :: * -> *) seq.
(Monad m, IsSequence seq, Element seq ~ Word8) =>
ConduitT seq seq m ()
linesUnboundedAsciiC ConduitT ByteString ByteString IO ()
-> ConduitM ByteString Void IO () -> ConduitM ByteString Void IO ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| (ByteString -> ConduitM ByteString Void IO ())
-> ConduitM ByteString Void IO ()
forall (m :: * -> *) i o r.
Monad m =>
(i -> ConduitT i o m r) -> ConduitT i o m ()
awaitForever (IO () -> ConduitM ByteString Void IO ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ConduitM ByteString Void IO ())
-> (ByteString -> IO ())
-> ByteString
-> ConduitM ByteString Void IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m () -> IO ()
forall a. m a -> IO a
unlift (m () -> IO ()) -> (ByteString -> m ()) -> ByteString -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> m ()
stderrLineHandler Int
pid))
      let eventConduit :: ConduitM a event m ()
eventConduit = Handle -> ConduitT a ByteString m ()
forall (m :: * -> *) i.
MonadIO m =>
Handle -> ConduitT i ByteString m ()
sourceHandle Handle
outHandle ConduitT a ByteString m ()
-> ConduitM ByteString event m () -> ConduitM a event m ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitM ByteString event m ()
forall b (m :: * -> *).
(Binary b, MonadThrow m) =>
ConduitT ByteString b m ()
conduitDecode
          commandConduit :: ConduitM a c IO ()
commandConduit =
            Chan (Maybe command) -> ConduitT a command IO ()
forall (m :: * -> *) a i.
MonadIO m =>
Chan (Maybe a) -> ConduitT i a m ()
sourceChan Chan (Maybe command)
commandChan
              ConduitT a command IO ()
-> ConduitM command c IO () -> ConduitM a c IO ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitT command ByteString IO ()
forall b (m :: * -> *).
(Binary b, MonadThrow m) =>
ConduitT b ByteString m ()
conduitEncode
              ConduitT command ByteString IO ()
-> ConduitM ByteString c IO () -> ConduitM command c IO ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| (ByteString -> [Flush ByteString])
-> ConduitT ByteString (Element [Flush ByteString]) IO ()
forall (m :: * -> *) mono a.
(Monad m, MonoFoldable mono) =>
(a -> mono) -> ConduitT a (Element mono) m ()
concatMapC (\ByteString
x -> [ByteString -> Flush ByteString
forall a. a -> Flush a
Chunk ByteString
x, Flush ByteString
forall a. Flush a
Flush])
              ConduitT ByteString (Flush ByteString) IO ()
-> ConduitM (Flush ByteString) c IO ()
-> ConduitM ByteString c IO ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| (IOError -> ConduitM (Flush ByteString) c IO ())
-> ConduitM (Flush ByteString) c IO ()
-> ConduitM (Flush ByteString) c IO ()
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 IOError -> ConduitM (Flush ByteString) c IO ()
forall {f :: * -> *}. MonadIO f => IOError -> f ()
handleEPIPE (Handle -> ConduitM (Flush ByteString) c IO ()
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 IOErrorType -> IOErrorType -> Bool
forall a. Eq a => a -> a -> Bool
== IOErrorType
ResourceVanished = () -> f ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          handleEPIPE IOError
e = IOError -> f ()
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO IOError
e
      let cmdThread :: IO ()
cmdThread = ConduitT () Void IO () -> IO ()
forall (m :: * -> *) r. Monad m => ConduitT () Void m r -> m r
runConduit ConduitT () Void IO ()
forall {a} {c}. ConduitM a c IO ()
commandConduit IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally` Handle -> IO ()
hClose Handle
outHandle
          eventThread :: IO ()
eventThread = m () -> IO ()
forall a. m a -> IO a
unlift (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ ConduitT () event m () -> (event -> m ()) -> m ()
forall (m :: * -> *) o a.
(MonadUnliftIO m, MonadIO m) =>
ConduitT () o m a -> (o -> m ()) -> m a
conduitToCallbacks ConduitT () event m ()
forall {a}. ConduitM a event m ()
eventConduit event -> m ()
eventHandler
      -- plain forkIO so it can process all of stderr in case of an exception
      IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO IO ()
stderrPiper
      IO ExitCode -> (Async ExitCode -> IO ExitCode) -> IO ExitCode
forall a b. IO a -> (Async a -> IO b) -> IO b
withAsync (ProcessHandle -> IO ExitCode
waitForProcess ProcessHandle
processHandle) ((Async ExitCode -> IO ExitCode) -> IO ExitCode)
-> (Async ExitCode -> IO ExitCode) -> IO ExitCode
forall a b. (a -> b) -> a -> b
$ \Async ExitCode
exitAsync -> do
        IO () -> (Async () -> IO ExitCode) -> IO ExitCode
forall a b. IO a -> (Async a -> IO b) -> IO b
withAsync IO ()
cmdThread ((Async () -> IO ExitCode) -> IO ExitCode)
-> (Async () -> IO ExitCode) -> IO ExitCode
forall a b. (a -> b) -> a -> b
$ \Async ()
_ -> do
          IO ()
eventThread
            IO () -> (SomeException -> IO ()) -> IO ()
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000
              Maybe ExitCode
maybeStatus <- Int -> IO ExitCode -> IO (Maybe ExitCode)
forall a. Int -> IO a -> IO (Maybe a)
timeout (Int
5 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
oneSecond) (Async ExitCode -> IO ExitCode
forall a. Async a -> IO a
wait Async ExitCode
exitAsync)
              WorkerException -> IO ()
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO (WorkerException -> IO ()) -> WorkerException -> IO ()
forall a b. (a -> b) -> a -> b
$ SomeException -> Maybe ExitCode -> WorkerException
WorkerException SomeException
e Maybe ExitCode
maybeStatus
          Async ExitCode -> IO ExitCode
forall a. Async a -> IO a
wait Async ExitCode
exitAsync