-- | This module defines the code for actually executing a command with tracing
-- enabled.
module HotelCalifornia.Exec where

import qualified Control.Exception as Exception
import qualified Data.Char as Char
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NEL
import Data.Maybe (fromMaybe)
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Text as Text
import Data.Text (Text)
import HotelCalifornia.Tracing
import HotelCalifornia.Tracing.TraceParent
import qualified OpenTelemetry.Trace.Core as Otel
import OpenTelemetry.Trace (Attribute(..), PrimitiveAttribute(..))
import Options.Applicative hiding (command)
import System.Environment (getEnvironment)
import System.Exit
import qualified System.Posix.Escape.Unicode as Escape
import System.Process.Typed
import HotelCalifornia.Which (which)

data Subprocess = Proc (NonEmpty String) | Shell String

commandToString :: Subprocess -> String
commandToString :: Subprocess -> String
commandToString (Proc NonEmpty String
tokens) = [String] -> String
Escape.escapeMany (NonEmpty String -> [String]
forall a. NonEmpty a -> [a]
NEL.toList NonEmpty String
tokens)
commandToString (Shell String
line) = String
line

commandToProcessConfig :: Subprocess -> ProcessConfig () () ()
commandToProcessConfig :: Subprocess -> ProcessConfig () () ()
commandToProcessConfig (Proc (String
command :| [String]
args)) = String -> [String] -> ProcessConfig () () ()
proc String
command [String]
args
commandToProcessConfig (Shell String
line) = String -> ProcessConfig () () ()
shell String
line

data ExecArgs = ExecArgs
    { ExecArgs -> Subprocess
execArgsSubprocess :: Subprocess
    , ExecArgs -> Maybe Text
execArgsSpanName :: Maybe Text
    , ExecArgs -> SpanStatus'
execArgsSigintStatus :: SpanStatus'
    , ExecArgs -> HashMap Text Attribute
execArgsAttributes :: HashMap Text Attribute
    }

-- | A variant of 'SpanStatus' that does not include a 'Text' for error.
data SpanStatus'
    = SpanUnset
    | SpanOk
    | SpanError

parseSpanStatus' :: ReadM SpanStatus'
parseSpanStatus' :: ReadM SpanStatus'
parseSpanStatus' = (String -> Either String SpanStatus') -> ReadM SpanStatus'
forall a. (String -> Either String a) -> ReadM a
eitherReader \String
s ->
    case (Char -> Char) -> String -> String
forall a b. (a -> b) -> [a] -> [b]
map Char -> Char
Char.toLower String
s of
        String
"unset" -> SpanStatus' -> Either String SpanStatus'
forall a b. b -> Either a b
Right SpanStatus'
SpanUnset
        String
"ok" -> SpanStatus' -> Either String SpanStatus'
forall a b. b -> Either a b
Right SpanStatus'
SpanOk
        String
"error" -> SpanStatus' -> Either String SpanStatus'
forall a b. b -> Either a b
Right SpanStatus'
SpanError
        String
_ -> String -> Either String SpanStatus'
forall a b. a -> Either a b
Left (String -> Either String SpanStatus')
-> String -> Either String SpanStatus'
forall a b. (a -> b) -> a -> b
$ [String] -> String
forall a. Monoid a => [a] -> a
mconcat [String
"Expected one of `unset`, `ok`, or `error` for SPAN_STATUS. Got: ", String
s]

parseProc :: Parser (NonEmpty String)
parseProc :: Parser (NonEmpty String)
parseProc = do
    String
command <- ReadM String -> Mod ArgumentFields String -> Parser String
forall a. ReadM a -> Mod ArgumentFields a -> Parser a
argument ReadM String
forall s. IsString s => ReadM s
str (String -> Mod ArgumentFields String
forall (f :: * -> *) a. HasMetavar f => String -> Mod f a
metavar String
"COMMAND")
    [String]
arguments <- Parser String -> Parser [String]
forall a. Parser a -> Parser [a]
forall (f :: * -> *) a. Alternative f => f a -> f [a]
many (ReadM String -> Mod ArgumentFields String -> Parser String
forall a. ReadM a -> Mod ArgumentFields a -> Parser a
argument ReadM String
forall s. IsString s => ReadM s
str (String -> Mod ArgumentFields String
forall (f :: * -> *) a. HasMetavar f => String -> Mod f a
metavar String
"ARGUMENT"))
    return (String
command String -> [String] -> NonEmpty String
forall a. a -> [a] -> NonEmpty a
:| [String]
arguments)

parseShell :: Parser String
parseShell :: Parser String
parseShell =
    ReadM String -> Mod OptionFields String -> Parser String
forall a. ReadM a -> Mod OptionFields a -> Parser a
option ReadM String
forall s. IsString s => ReadM s
str
        (   String -> Mod OptionFields String
forall (f :: * -> *) a. HasMetavar f => String -> Mod f a
metavar String
"SCRIPT"
        Mod OptionFields String
-> Mod OptionFields String -> Mod OptionFields String
forall a. Semigroup a => a -> a -> a
<>  String -> Mod OptionFields String
forall (f :: * -> *) a. HasName f => String -> Mod f a
long String
"shell"
        Mod OptionFields String
-> Mod OptionFields String -> Mod OptionFields String
forall a. Semigroup a => a -> a -> a
<>  String -> Mod OptionFields String
forall (f :: * -> *) a. String -> Mod f a
help String
"Run an arbitrary shell script instead of running an executable command"
        )

parseSubprocess :: Parser Subprocess
parseSubprocess :: Parser Subprocess
parseSubprocess = (NonEmpty String -> Subprocess)
-> Parser (NonEmpty String) -> Parser Subprocess
forall a b. (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NonEmpty String -> Subprocess
Proc Parser (NonEmpty String)
parseProc Parser Subprocess -> Parser Subprocess -> Parser Subprocess
forall a. Parser a -> Parser a -> Parser a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (String -> Subprocess) -> Parser String -> Parser Subprocess
forall a b. (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap String -> Subprocess
Shell Parser String
parseShell

-- | Parse a `key=value` string into an attribute.
parseAttribute :: String -> Either String (Text, Attribute)
parseAttribute :: String -> Either String (Text, Attribute)
parseAttribute String
input = do
    let (Text
key, Text
value') = HasCallStack => Text -> Text -> (Text, Text)
Text -> Text -> (Text, Text)
Text.breakOn Text
"=" (Text -> (Text, Text)) -> Text -> (Text, Text)
forall a b. (a -> b) -> a -> b
$ String -> Text
Text.pack String
input
    if Text -> Bool
Text.null Text
value' Bool -> Bool -> Bool
|| Text -> Bool
Text.null Text
key
    then String -> Either String (Text, Attribute)
forall a b. a -> Either a b
Left (String -> Either String (Text, Attribute))
-> String -> Either String (Text, Attribute)
forall a b. (a -> b) -> a -> b
$ String
"Attributes must contain a non-empty key and value separated by `=`: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
input
    else (Text, Attribute) -> Either String (Text, Attribute)
forall a. a -> Either String a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Text, Attribute) -> Either String (Text, Attribute))
-> (Text, Attribute) -> Either String (Text, Attribute)
forall a b. (a -> b) -> a -> b
$ (Text
key, PrimitiveAttribute -> Attribute
AttributeValue (PrimitiveAttribute -> Attribute)
-> PrimitiveAttribute -> Attribute
forall a b. (a -> b) -> a -> b
$ Text -> PrimitiveAttribute
TextAttribute (Text -> PrimitiveAttribute) -> Text -> PrimitiveAttribute
forall a b. (a -> b) -> a -> b
$ Int -> Text -> Text
Text.drop Int
1 Text
value')

parseExecArgs :: Parser ExecArgs
parseExecArgs :: Parser ExecArgs
parseExecArgs = do
    Maybe Text
execArgsSpanName <- Parser Text -> Parser (Maybe Text)
forall (f :: * -> *) a. Alternative f => f a -> f (Maybe a)
optional do
        ReadM Text -> Mod OptionFields Text -> Parser Text
forall a. ReadM a -> Mod OptionFields a -> Parser a
option ReadM Text
forall s. IsString s => ReadM s
str (Mod OptionFields Text -> Parser Text)
-> Mod OptionFields Text -> Parser Text
forall a b. (a -> b) -> a -> b
$ [Mod OptionFields Text] -> Mod OptionFields Text
forall a. Monoid a => [a] -> a
mconcat
            [ String -> Mod OptionFields Text
forall (f :: * -> *) a. HasMetavar f => String -> Mod f a
metavar String
"SPAN_NAME"
            , String -> Mod OptionFields Text
forall (f :: * -> *) a. HasName f => String -> Mod f a
long String
"span-name"
            , Char -> Mod OptionFields Text
forall (f :: * -> *) a. HasName f => Char -> Mod f a
short Char
's'
            , String -> Mod OptionFields Text
forall (f :: * -> *) a. String -> Mod f a
help String
"The name of the span that the program reports. By default, this is the script you pass in."
            ]
    SpanStatus'
execArgsSigintStatus <-
        ReadM SpanStatus'
-> Mod OptionFields SpanStatus' -> Parser SpanStatus'
forall a. ReadM a -> Mod OptionFields a -> Parser a
option ReadM SpanStatus'
parseSpanStatus' (Mod OptionFields SpanStatus' -> Parser SpanStatus')
-> Mod OptionFields SpanStatus' -> Parser SpanStatus'
forall a b. (a -> b) -> a -> b
$ [Mod OptionFields SpanStatus'] -> Mod OptionFields SpanStatus'
forall a. Monoid a => [a] -> a
mconcat
            [ String -> Mod OptionFields SpanStatus'
forall (f :: * -> *) a. HasMetavar f => String -> Mod f a
metavar String
"SPAN_STATUS"
            , String -> Mod OptionFields SpanStatus'
forall (f :: * -> *) a. HasName f => String -> Mod f a
long String
"set-sigint-status"
            , Char -> Mod OptionFields SpanStatus'
forall (f :: * -> *) a. HasName f => Char -> Mod f a
short Char
'i'
            , String -> Mod OptionFields SpanStatus'
forall (f :: * -> *) a. String -> Mod f a
help String
"The status reported when the process is killed with SIGINT."
            , SpanStatus' -> Mod OptionFields SpanStatus'
forall (f :: * -> *) a. HasValue f => a -> Mod f a
value SpanStatus'
SpanUnset
            ]
    HashMap Text Attribute
execArgsAttributes <-
        [(Text, Attribute)] -> HashMap Text Attribute
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList ([(Text, Attribute)] -> HashMap Text Attribute)
-> Parser [(Text, Attribute)] -> Parser (HashMap Text Attribute)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser (Text, Attribute) -> Parser [(Text, Attribute)]
forall a. Parser a -> Parser [a]
forall (f :: * -> *) a. Alternative f => f a -> f [a]
many (ReadM (Text, Attribute)
-> Mod OptionFields (Text, Attribute) -> Parser (Text, Attribute)
forall a. ReadM a -> Mod OptionFields a -> Parser a
option ((String -> Either String (Text, Attribute))
-> ReadM (Text, Attribute)
forall a. (String -> Either String a) -> ReadM a
eitherReader String -> Either String (Text, Attribute)
parseAttribute) (Mod OptionFields (Text, Attribute) -> Parser (Text, Attribute))
-> Mod OptionFields (Text, Attribute) -> Parser (Text, Attribute)
forall a b. (a -> b) -> a -> b
$ [Mod OptionFields (Text, Attribute)]
-> Mod OptionFields (Text, Attribute)
forall a. Monoid a => [a] -> a
mconcat
            [ String -> Mod OptionFields (Text, Attribute)
forall (f :: * -> *) a. HasMetavar f => String -> Mod f a
metavar String
"KEY=VALUE"
            , String -> Mod OptionFields (Text, Attribute)
forall (f :: * -> *) a. HasName f => String -> Mod f a
long String
"attribute"
            , Char -> Mod OptionFields (Text, Attribute)
forall (f :: * -> *) a. HasName f => Char -> Mod f a
short Char
'a'
            , String -> Mod OptionFields (Text, Attribute)
forall (f :: * -> *) a. String -> Mod f a
help String
"A string attribute to add to the span."
            ])
    Subprocess
execArgsSubprocess <- Parser Subprocess
parseSubprocess
    pure ExecArgs{Maybe Text
HashMap Text Attribute
SpanStatus'
Subprocess
execArgsSubprocess :: Subprocess
execArgsSpanName :: Maybe Text
execArgsSigintStatus :: SpanStatus'
execArgsAttributes :: HashMap Text Attribute
execArgsSpanName :: Maybe Text
execArgsSigintStatus :: SpanStatus'
execArgsAttributes :: HashMap Text Attribute
execArgsSubprocess :: Subprocess
..}

makeInitialAttributes :: Subprocess -> HashMap Text Attribute -> IO (HashMap Text Attribute)
makeInitialAttributes :: Subprocess -> HashMap Text Attribute -> IO (HashMap Text Attribute)
makeInitialAttributes Subprocess
subprocess HashMap Text Attribute
extraAttributes = do
    HashMap Text Attribute
processAttributes <-
        case Subprocess
subprocess of
          Proc (String
command :| [String]
args) -> do
              [(Text, Attribute)]
pathAttribute <-
                  ((String -> [(Text, Attribute)])
-> Maybe String -> [(Text, Attribute)]
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\String
path -> [(Text
executablePathName, Text -> Attribute
forall a. ToAttribute a => a -> Attribute
Otel.toAttribute (Text -> Attribute) -> Text -> Attribute
forall a b. (a -> b) -> a -> b
$ String -> Text
Text.pack String
path)]))
                  (Maybe String -> [(Text, Attribute)])
-> IO (Maybe String) -> IO [(Text, Attribute)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Maybe String)
which String
command
              pure $ [(Text, Attribute)] -> HashMap Text Attribute
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList ([(Text, Attribute)] -> HashMap Text Attribute)
-> [(Text, Attribute)] -> HashMap Text Attribute
forall a b. (a -> b) -> a -> b
$
                [ (Text
commandArgsName, [Text] -> Attribute
forall a. ToAttribute a => a -> Attribute
Otel.toAttribute ([Text] -> Attribute) -> [Text] -> Attribute
forall a b. (a -> b) -> a -> b
$ (String -> Text) -> [String] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map String -> Text
Text.pack [String]
args)
                , (Text
executableName, Text -> Attribute
forall a. ToAttribute a => a -> Attribute
Otel.toAttribute (Text -> Attribute) -> Text -> Attribute
forall a b. (a -> b) -> a -> b
$ String -> Text
Text.pack String
command)
                ] [(Text, Attribute)] -> [(Text, Attribute)] -> [(Text, Attribute)]
forall a. Semigroup a => a -> a -> a
<> [(Text, Attribute)]
pathAttribute
          Shell String
_command -> HashMap Text Attribute -> IO (HashMap Text Attribute)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HashMap Text Attribute
forall a. Monoid a => a
mempty

    pure $ HashMap Text Attribute
processAttributes HashMap Text Attribute
-> HashMap Text Attribute -> HashMap Text Attribute
forall a. Semigroup a => a -> a -> a
<> HashMap Text Attribute
extraAttributes

runNoTracing :: Subprocess -> IO ()
runNoTracing :: Subprocess -> IO ()
runNoTracing Subprocess
subproc = do
    let processConfig :: ProcessConfig () () ()
processConfig = Subprocess -> ProcessConfig () () ()
commandToProcessConfig Subprocess
subproc
    [(String, String)]
userEnv <- IO [(String, String)]
getEnvironment
    ExitCode
exitCode <- ProcessConfig () () () -> IO ExitCode
forall (m :: * -> *) stdin stdout stderr.
MonadIO m =>
ProcessConfig stdin stdout stderr -> m ExitCode
runProcess (ProcessConfig () () () -> IO ExitCode)
-> ProcessConfig () () () -> IO ExitCode
forall a b. (a -> b) -> a -> b
$ [(String, String)]
-> ProcessConfig () () () -> ProcessConfig () () ()
forall stdin stdout stderr.
[(String, String)]
-> ProcessConfig stdin stdout stderr
-> ProcessConfig stdin stdout stderr
setEnv [(String, String)]
userEnv ProcessConfig () () ()
processConfig
    ExitCode -> IO ()
forall a. ExitCode -> IO a
exitWith ExitCode
exitCode

runExecArgs :: ExecArgs -> IO ()
runExecArgs :: ExecArgs -> IO ()
runExecArgs ExecArgs {Maybe Text
HashMap Text Attribute
SpanStatus'
Subprocess
execArgsSubprocess :: ExecArgs -> Subprocess
execArgsSpanName :: ExecArgs -> Maybe Text
execArgsSigintStatus :: ExecArgs -> SpanStatus'
execArgsAttributes :: ExecArgs -> HashMap Text Attribute
execArgsSubprocess :: Subprocess
execArgsSpanName :: Maybe Text
execArgsSigintStatus :: SpanStatus'
execArgsAttributes :: HashMap Text Attribute
..} = do
    HashMap Text Attribute
initialAttributes <- Subprocess -> HashMap Text Attribute -> IO (HashMap Text Attribute)
makeInitialAttributes Subprocess
execArgsSubprocess HashMap Text Attribute
execArgsAttributes

    let script :: String
script = Subprocess -> String
commandToString Subprocess
execArgsSubprocess
        spanName :: Text
spanName =
            Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe (String -> Text
Text.pack String
script) Maybe Text
execArgsSpanName
        spanArguments :: SpanArguments
spanArguments = SpanArguments
defaultSpanArguments { Otel.attributes = initialAttributes }

    Text -> SpanArguments -> (Span -> IO ()) -> IO ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
Text -> SpanArguments -> (Span -> m a) -> m a
inSpanWith' Text
spanName SpanArguments
spanArguments \Span
span' -> do
        [(String, String)]
newEnv <- Span -> IO [(String, String)]
spanContextToEnvironment Span
span'
        [(String, String)]
fullEnv <- [(String, String)] -> [(String, String)] -> [(String, String)]
forall a. Monoid a => a -> a -> a
mappend [(String, String)]
newEnv ([(String, String)] -> [(String, String)])
-> IO [(String, String)] -> IO [(String, String)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO [(String, String)]
getEnvironment

        let processConfig :: ProcessConfig () () ()
processConfig = Subprocess -> ProcessConfig () () ()
commandToProcessConfig Subprocess
execArgsSubprocess

        let handleSigInt :: AsyncException -> IO (Maybe ExitCode)
handleSigInt =
                \case
                    AsyncException
Exception.UserInterrupt -> do
                        Span -> Text -> Int -> IO ()
forall (m :: * -> *) a.
(MonadIO m, ToAttribute a) =>
Span -> Text -> a -> m ()
Otel.addAttribute Span
span' Text
exitStatusName (Int
-2 :: Int) -- SIGINT
                        case SpanStatus'
execArgsSigintStatus of
                            SpanStatus'
SpanUnset ->
                                Maybe ExitCode -> IO (Maybe ExitCode)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ExitCode
forall a. Maybe a
Nothing
                            SpanStatus'
SpanOk -> do
                                Span -> SpanStatus -> IO ()
forall (m :: * -> *). MonadIO m => Span -> SpanStatus -> m ()
Otel.setStatus Span
span' SpanStatus
Otel.Ok
                                pure Maybe ExitCode
forall a. Maybe a
Nothing
                            SpanStatus'
SpanError -> do
                                -- `hs-opentelemetry` will automatically mark a
                                -- span as an error if it ends with an
                                -- exception.
                                AsyncException -> IO (Maybe ExitCode)
forall e a. Exception e => e -> IO a
Exception.throwIO AsyncException
Exception.UserInterrupt
                    AsyncException
other ->
                        AsyncException -> IO (Maybe ExitCode)
forall e a. Exception e => e -> IO a
Exception.throwIO AsyncException
other

        Maybe ExitCode
mexitCode <- (AsyncException -> IO (Maybe ExitCode))
-> IO (Maybe ExitCode) -> IO (Maybe ExitCode)
forall e a. Exception e => (e -> IO a) -> IO a -> IO a
Exception.handle AsyncException -> IO (Maybe ExitCode)
handleSigInt (IO (Maybe ExitCode) -> IO (Maybe ExitCode))
-> IO (Maybe ExitCode) -> IO (Maybe ExitCode)
forall a b. (a -> b) -> a -> b
$ (ExitCode -> Maybe ExitCode) -> IO ExitCode -> IO (Maybe ExitCode)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ExitCode -> Maybe ExitCode
forall a. a -> Maybe a
Just (IO ExitCode -> IO (Maybe ExitCode))
-> IO ExitCode -> IO (Maybe ExitCode)
forall a b. (a -> b) -> a -> b
$ ProcessConfig () () () -> IO ExitCode
forall (m :: * -> *) stdin stdout stderr.
MonadIO m =>
ProcessConfig stdin stdout stderr -> m ExitCode
runProcess (ProcessConfig () () () -> IO ExitCode)
-> ProcessConfig () () () -> IO ExitCode
forall a b. (a -> b) -> a -> b
$ [(String, String)]
-> ProcessConfig () () () -> ProcessConfig () () ()
forall stdin stdout stderr.
[(String, String)]
-> ProcessConfig stdin stdout stderr
-> ProcessConfig stdin stdout stderr
setEnv [(String, String)]
fullEnv ProcessConfig () () ()
processConfig

        case Maybe ExitCode
mexitCode of
            Just ExitCode
exitCode -> do
                Span -> Text -> Int -> IO ()
forall (m :: * -> *) a.
(MonadIO m, ToAttribute a) =>
Span -> Text -> a -> m ()
Otel.addAttribute Span
span' Text
exitStatusName
                    case ExitCode
exitCode of
                       ExitCode
ExitSuccess -> Int
0
                       ExitFailure Int
status -> Int
status
                case ExitCode
exitCode of
                    ExitCode
ExitSuccess -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                    ExitFailure Int
_ -> ExitCode -> IO ()
forall a. ExitCode -> IO a
exitWith ExitCode
exitCode
            Maybe ExitCode
Nothing ->
                () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

exitStatusName :: Text
exitStatusName :: Text
exitStatusName = Text
"process.exit_status"

executablePathName :: Text
executablePathName :: Text
executablePathName = Text
"process.executable.path"

executableName :: Text
executableName :: Text
executableName = Text
"process.executable.path"

commandArgsName :: Text
commandArgsName :: Text
commandArgsName = Text
"process.command_args"