Exec.hs

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

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

data Subprocess = Proc (NonEmpty String) | Shell String

commandToString :: Subprocess -> String
commandToString (Proc tokens) = Escape.escapeMany (NEL.toList tokens)
commandToString (Shell line) = line

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

data ExecArgs = ExecArgs
    { execArgsSubprocess :: Subprocess
    , execArgsSpanName :: Maybe Text
    , execArgsSigintStatus :: SpanStatus'
    , 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' = eitherReader \s ->
    case map Char.toLower s of
        "unset" -> Right SpanUnset
        "ok" -> Right SpanOk
        "error" -> Right SpanError
        _ ->
            Left $
                mconcat ["Expected one of `unset`, `ok`, or `error` for SPAN_STATUS. Got: ", s]

parseProc :: Parser (NonEmpty String)
parseProc = do
    command <- argument str (metavar "COMMAND")
    arguments <- many (argument str (metavar "ARGUMENT"))
    return (command :| arguments)

parseShell :: Parser String
parseShell =
    option
        str
        ( metavar "SCRIPT"
            <> long "shell"
            <> help "Run an arbitrary shell script instead of running an executable command"
        )

parseSubprocess :: Parser Subprocess
parseSubprocess = fmap Proc parseProc <|> fmap Shell parseShell

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

parseExecArgs :: Parser ExecArgs
parseExecArgs = do
    execArgsSpanName <- optional do
        option str $
            mconcat
                [ metavar "SPAN_NAME"
                , long "span-name"
                , short 's'
                , help
                    "The name of the span that the program reports. By default, this is the script you pass in."
                ]
    execArgsSigintStatus <-
        option parseSpanStatus' $
            mconcat
                [ metavar "SPAN_STATUS"
                , long "set-sigint-status"
                , short 'i'
                , help "The status reported when the process is killed with SIGINT."
                , value SpanUnset
                ]
    execArgsAttributes <-
        HashMap.fromList
            <$> many
                ( option (eitherReader parseAttribute) $
                    mconcat
                        [ metavar "KEY=VALUE"
                        , long "attribute"
                        , short 'a'
                        , help "A string attribute to add to the span."
                        ]
                )
    execArgsSubprocess <- parseSubprocess
    pure ExecArgs{..}

makeInitialAttributes
    :: Subprocess -> HashMap Text Attribute -> IO (HashMap Text Attribute)
makeInitialAttributes subprocess extraAttributes = do
    processAttributes <-
        case subprocess of
            Proc (command :| args) -> do
                pathAttribute <-
                    (foldMap (\path -> [(executablePathName, Otel.toAttribute $ Text.pack path)]))
                        <$> which command
                pure $
                    HashMap.fromList $
                        [ (commandArgsName, Otel.toAttribute $ map Text.pack args)
                        , (executableName, Otel.toAttribute $ Text.pack command)
                        ]
                            <> pathAttribute
            Shell _command -> pure mempty

    pure $ processAttributes <> extraAttributes

runNoTracing :: Subprocess -> IO ()
runNoTracing subproc = do
    let
        processConfig = commandToProcessConfig subproc
    userEnv <- getEnvironment
    exitCode <- runProcess $ setEnv userEnv processConfig
    exitWith exitCode

runExecArgs :: ExecArgs -> IO ()
runExecArgs ExecArgs{..} = do
    initialAttributes <- makeInitialAttributes execArgsSubprocess execArgsAttributes

    let
        script = commandToString execArgsSubprocess
        spanName =
            fromMaybe (Text.pack script) execArgsSpanName
        spanArguments = defaultSpanArguments{Otel.attributes = initialAttributes}

    inSpanWith' spanName spanArguments \span' -> do
        newEnv <- spanContextToEnvironment span'
        fullEnv <- mappend newEnv <$> getEnvironment

        let
            processConfig = commandToProcessConfig execArgsSubprocess

        let
            handleSigInt =
                \case
                    Exception.UserInterrupt -> do
                        Otel.addAttribute span' exitStatusName (-2 :: Int) -- SIGINT
                        case execArgsSigintStatus of
                            SpanUnset ->
                                pure Nothing
                            SpanOk -> do
                                Otel.setStatus span' Otel.Ok
                                pure Nothing
                            SpanError -> do
                                -- `hs-opentelemetry` will automatically mark a
                                -- span as an error if it ends with an
                                -- exception.
                                Exception.throwIO Exception.UserInterrupt
                    other ->
                        Exception.throwIO other

        mexitCode <-
            Exception.handle handleSigInt $
                fmap Just $
                    runProcess $
                        setEnv fullEnv processConfig

        case mexitCode of
            Just exitCode -> do
                Otel.addAttribute
                    span'
                    exitStatusName
                    case exitCode of
                        ExitSuccess -> 0
                        ExitFailure status -> status
                case exitCode of
                    ExitSuccess -> pure ()
                    ExitFailure _ -> exitWith exitCode
            Nothing ->
                pure ()

exitStatusName :: Text
exitStatusName = "process.exit_status"

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

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

commandArgsName :: Text
commandArgsName = "process.command_args"