{-# LANGUAGE BlockArguments    #-}
{-# LANGUAGE NamedFieldPuns    #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE GADTs #-}

{-| This module provides a GHC plugin that will export open telemetry metrics
    for your build.  Specifically, this plugin will create one span per module
    (recording how long that module took to build) and one sub-span per phase
    of that module's build (recording how long that phase took).
-}
module OpenTelemetry.Plugin
    ( -- * Plugin
      plugin
    ) where

import Control.Monad.IO.Class (MonadIO(..))
import GHC.Types.Target (Target(..), TargetId(..))
import OpenTelemetry.Context (Context)

import GHC.Driver.Pipeline (runPhase)
import GHC.Driver.Pipeline.Phases
  ( PhaseHook (..),
    TPhase (..),
  )
import GHC.Driver.Hooks (Hooks (..))
import GHC.Plugins
    ( CoreToDo(..)
    , HscEnv(..)
    , Plugin(..)
    )
import qualified Data.Text as Text
import qualified GHC.Plugins as Plugins
import qualified GHC.Utils.Outputable as Outputable
import qualified OpenTelemetry.Plugin.Shared as Shared
import qualified GHC.Driver.Backend as Backend

wrapTodo :: MonadIO io => IO Context -> CoreToDo -> io CoreToDo
wrapTodo :: forall (io :: * -> *).
MonadIO io =>
IO Context -> CoreToDo -> io CoreToDo
wrapTodo IO Context
getParentContext CoreToDo
todo =
    case CoreToDo
todo of
        CoreDoPasses [CoreToDo]
passes ->
            ([CoreToDo] -> CoreToDo) -> io [CoreToDo] -> io CoreToDo
forall a b. (a -> b) -> io a -> io b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [CoreToDo] -> CoreToDo
CoreDoPasses ((CoreToDo -> io CoreToDo) -> [CoreToDo] -> io [CoreToDo]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (IO Context -> CoreToDo -> io CoreToDo
forall (io :: * -> *).
MonadIO io =>
IO Context -> CoreToDo -> io CoreToDo
wrapTodo IO Context
getParentContext) [CoreToDo]
passes)

        CoreToDo
_ -> IO CoreToDo -> io CoreToDo
forall a. IO a -> io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
            let sdoc :: SDoc
sdoc = CoreToDo -> SDoc
forall a. Outputable a => a -> SDoc
Outputable.ppr CoreToDo
todo

            let label :: String
label =
                    SDocContext -> SDoc -> String
Outputable.showSDocOneLine SDocContext
Outputable.defaultSDocContext SDoc
sdoc

            (IO Context
_, IO ()
beginPass, IO ()
endPass) <- do
                Bool -> IO Context -> Text -> IO (IO Context, IO (), IO ())
Shared.makeWrapperPluginPasses Bool
False IO Context
getParentContext (String -> Text
Text.pack String
label)

            let beginPluginPass :: CoreToDo
beginPluginPass =
                    String -> CorePluginPass -> CoreToDo
CoreDoPluginPass (String
"begin " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
label) \ModGuts
modGuts -> IO ModGuts -> CoreM ModGuts
forall a. IO a -> CoreM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
                        IO ()
beginPass

                        ModGuts -> IO ModGuts
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ModGuts
modGuts

            let endPluginPass :: CoreToDo
endPluginPass =
                    String -> CorePluginPass -> CoreToDo
CoreDoPluginPass (String
"end " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
label) \ModGuts
modGuts -> IO ModGuts -> CoreM ModGuts
forall a. IO a -> CoreM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
                        IO ()
endPass

                        ModGuts -> IO ModGuts
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ModGuts
modGuts

            CoreToDo -> IO CoreToDo
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([CoreToDo] -> CoreToDo
CoreDoPasses [ CoreToDo
beginPluginPass, CoreToDo
todo, CoreToDo
endPluginPass ])

-- | GHC plugin that exports open telemetry metrics about the build
plugin :: Plugin
plugin :: Plugin
plugin =
    Plugin
Plugins.defaultPlugin
        { driverPlugin
        , pluginRecompile
        , installCoreToDos
        }
  where
    driverPlugin :: p -> HscEnv -> IO HscEnv
driverPlugin p
_ hscEnv :: HscEnv
hscEnv@HscEnv{ [Target]
hsc_targets :: [Target]
hsc_targets :: HscEnv -> [Target]
hsc_targets } = do
        let rootModuleNames :: [String]
rootModuleNames = do
                Target{ targetId :: Target -> TargetId
targetId = TargetModule ModuleName
rootModuleName } <- [Target]
hsc_targets

                String -> [String]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ModuleName -> String
Plugins.moduleNameString ModuleName
rootModuleName)

        let closePhase :: ClosePhase
closePhase =
                case Backend -> Bool
Backend.backendWritesFiles (Backend -> Bool) -> Backend -> Bool
forall a b. (a -> b) -> a -> b
$ DynFlags -> Backend
Plugins.backend (DynFlags -> Backend) -> DynFlags -> Backend
forall a b. (a -> b) -> a -> b
$ HscEnv -> DynFlags
hsc_dflags HscEnv
hscEnv of
                    Bool
False ->
                        ClosePhase
CloseInHscBackend
                    Bool
True ->
                        ClosePhase
CloseInMergeForeign

        [String] -> IO ()
Shared.setRootModuleNames [String]
rootModuleNames

        let packageName :: PackageName
packageName = HscEnv -> PackageName
getPackageName HscEnv
hscEnv

        PackageName -> IO ()
Shared.initializeTopLevelContext PackageName
packageName

        HscEnv -> IO HscEnv
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HscEnv
hscEnv
            { hsc_hooks =
                (hsc_hooks hscEnv)
                    { runPhaseHook =
                        Just $ PhaseHook \TPhase a
phase -> do
                            case TPhase a
phase of
                                T_Hsc HscEnv
_ ModSummary
modSummary -> do
                                    let modName :: String
modName =
                                            ModuleName -> String
Plugins.moduleNameString (ModuleName -> String)
-> (ModSummary -> ModuleName) -> ModSummary -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenModule Unit -> ModuleName
forall unit. GenModule unit -> ModuleName
Plugins.moduleName (GenModule Unit -> ModuleName)
-> (ModSummary -> GenModule Unit) -> ModSummary -> ModuleName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ModSummary -> GenModule Unit
Plugins.ms_mod (ModSummary -> String) -> ModSummary -> String
forall a b. (a -> b) -> a -> b
$
                                                ModSummary
modSummary
                                        modObjectLocation :: String
modObjectLocation =
                                            ModLocation -> String
Plugins.ml_obj_file (ModLocation -> String) -> ModLocation -> String
forall a b. (a -> b) -> a -> b
$ ModSummary -> ModLocation
Plugins.ms_location ModSummary
modSummary
                                    PackageName -> String -> String -> IO ()
Shared.recordModuleStart PackageName
packageName String
modObjectLocation String
modName
                                    TPhase a -> IO a
forall a. TPhase a -> IO a
runPhase TPhase a
phase
                                T_MergeForeign PipeEnv
_pipeEnv HscEnv
_hscEnv String
objectFilePath [String]
_filePaths -> do
                                    -- this phase appears to only be run
                                    -- during compilation, not ghci
                                    a
x <- TPhase a -> IO a
forall a. TPhase a -> IO a
runPhase TPhase a
phase
                                    case ClosePhase
closePhase of
                                        ClosePhase
CloseInMergeForeign ->
                                            String -> IO ()
Shared.recordModuleEnd String
objectFilePath
                                        ClosePhase
_ ->
                                            () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                                    a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
                                T_HscBackend PipeEnv
_pipeEnv HscEnv
_hscEnv ModuleName
modName HscSource
_hscSrc ModLocation
_modLoc HscBackendAction
_hscAction  -> do
                                    -- this happens in ghci for sure as
                                    -- a last step
                                    a
x <- TPhase a -> IO a
forall a. TPhase a -> IO a
runPhase TPhase a
phase
                                    case ClosePhase
closePhase of
                                        ClosePhase
CloseInHscBackend ->
                                            String -> IO ()
Shared.recordModuleEnd (ModuleName -> String
Plugins.moduleNameString ModuleName
modName)
                                        ClosePhase
_ ->
                                            () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                                    a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
                                TPhase a
_ -> do

                                    TPhase a -> IO a
forall a. TPhase a -> IO a
runPhase TPhase a
phase
                    }
            }

    installCoreToDos :: p -> [CoreToDo] -> m [CoreToDo]
installCoreToDos p
_ [CoreToDo]
todos = do
        Bool
shouldMakeSubPasses <- IO Bool -> m Bool
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Bool
Shared.getPluginShouldRecordPasses

        if Bool
shouldMakeSubPasses
            then do
                GenModule Unit
module_ <- m (GenModule Unit)
forall (m :: * -> *). HasModule m => m (GenModule Unit)
Plugins.getModule

                (IO Context
getCurrentContext, IO ()
firstPluginPass, IO ()
lastPluginPass) <- do
                    let moduleNameString :: String
moduleNameString =
                            ModuleName -> String
Plugins.moduleNameString (ModuleName -> String) -> ModuleName -> String
forall a b. (a -> b) -> a -> b
$ GenModule Unit -> ModuleName
forall unit. GenModule unit -> ModuleName
Plugins.moduleName GenModule Unit
module_
                        getContext :: IO Context
getContext =
                            String -> IO Context -> IO Context
Shared.modifyContextWithParentSpan String
moduleNameString IO Context
Shared.getTopLevelContext
                    IO (IO Context, IO (), IO ()) -> m (IO Context, IO (), IO ())
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Bool -> IO Context -> Text -> IO (IO Context, IO (), IO ())
Shared.makeWrapperPluginPasses Bool
True IO Context
getContext Text
"CoreToDos")

                let firstPass :: CoreToDo
firstPass =
                        String -> CorePluginPass -> CoreToDo
CoreDoPluginPass String
"begin module" \ModGuts
modGuts -> IO ModGuts -> CoreM ModGuts
forall a. IO a -> CoreM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
                            IO ()
firstPluginPass

                            ModGuts -> IO ModGuts
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ModGuts
modGuts

                let lastPass :: CoreToDo
lastPass =
                        String -> CorePluginPass -> CoreToDo
CoreDoPluginPass String
"end module" \ModGuts
modGuts -> IO ModGuts -> CoreM ModGuts
forall a. IO a -> CoreM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
                            IO ()
lastPluginPass

                            ModGuts -> IO ModGuts
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ModGuts
modGuts

                [CoreToDo]
newTodos <- (CoreToDo -> m CoreToDo) -> [CoreToDo] -> m [CoreToDo]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (IO Context -> CoreToDo -> m CoreToDo
forall (io :: * -> *).
MonadIO io =>
IO Context -> CoreToDo -> io CoreToDo
wrapTodo IO Context
getCurrentContext) [CoreToDo]
todos
                [CoreToDo] -> m [CoreToDo]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ CoreToDo
firstPass ] [CoreToDo] -> [CoreToDo] -> [CoreToDo]
forall a. Semigroup a => a -> a -> a
<> [CoreToDo]
newTodos [CoreToDo] -> [CoreToDo] -> [CoreToDo]
forall a. Semigroup a => a -> a -> a
<> [ CoreToDo
lastPass ])
            else do
                [CoreToDo] -> m [CoreToDo]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [CoreToDo]
todos

    pluginRecompile :: [String] -> IO PluginRecompile
pluginRecompile = [String] -> IO PluginRecompile
Plugins.purePlugin

data ClosePhase = CloseInHscBackend | CloseInMergeForeign

getPackageName :: HscEnv -> Shared.PackageName
getPackageName :: HscEnv -> PackageName
getPackageName =
    Text -> PackageName
Shared.PackageName
        (Text -> PackageName) -> (HscEnv -> Text) -> HscEnv -> PackageName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
Text.pack
        (String -> Text) -> (HscEnv -> String) -> HscEnv -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnitId -> String
Plugins.unitIdString
        (UnitId -> String) -> (HscEnv -> UnitId) -> HscEnv -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DynFlags -> UnitId
Plugins.homeUnitId_
        (DynFlags -> UnitId) -> (HscEnv -> DynFlags) -> HscEnv -> UnitId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HscEnv -> DynFlags
Plugins.hsc_dflags