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

{-| 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.Plugins
    ( CoreToDo(..)
    , GenModule(..)
    , HscEnv(..)
    , Plugin(..)
    )
import qualified Control.Monad as Monad
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

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 ->
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [CoreToDo] -> CoreToDo
CoreDoPasses (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (io :: * -> *).
MonadIO io =>
IO Context -> CoreToDo -> io CoreToDo
wrapTodo IO Context
getParentContext) [CoreToDo]
passes)

        CoreToDo
_ -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
            let sdoc :: SDoc
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 " forall a. Semigroup a => a -> a -> a
<> String
label) \ModGuts
modGuts -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
                        IO ()
beginPass

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

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

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

            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
        { forall {p}. p -> HscEnv -> IO HscEnv
driverPlugin :: [String] -> HscEnv -> IO HscEnv
driverPlugin :: forall {p}. p -> HscEnv -> IO HscEnv
driverPlugin
        , [String] -> IO PluginRecompile
pluginRecompile :: [String] -> IO PluginRecompile
pluginRecompile :: [String] -> IO PluginRecompile
pluginRecompile
        , forall {m :: * -> *} {p}.
(HasModule m, MonadIO m) =>
p -> [CoreToDo] -> m [CoreToDo]
installCoreToDos :: CorePlugin
installCoreToDos :: forall {m :: * -> *} {p}.
(HasModule m, MonadIO m) =>
p -> [CoreToDo] -> m [CoreToDo]
installCoreToDos
        }
  where
    driverPlugin :: p -> HscEnv -> IO HscEnv
driverPlugin p
_ hscEnv :: HscEnv
hscEnv@HscEnv{ [Target]
hsc_targets :: HscEnv -> [Target]
hsc_targets :: [Target]
hsc_targets } = do
        let rootModuleNames :: [String]
rootModuleNames = do
                Target{ targetId :: Target -> TargetId
targetId = TargetModule ModuleName
rootModuleName } <- [Target]
hsc_targets

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

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

        IO ()
Shared.initializeTopLevelContext

        forall (f :: * -> *) a. Applicative f => a -> f a
pure HscEnv
hscEnv

    installCoreToDos :: p -> [CoreToDo] -> m [CoreToDo]
installCoreToDos p
_ [CoreToDo]
todos = do
        Module
module_ <- forall (m :: * -> *). HasModule m => m Module
Plugins.getModule

        let moduleName_ :: ModuleName
moduleName_ = forall unit. GenModule unit -> ModuleName
moduleName Module
module_

        let moduleText :: Text
moduleText = String -> Text
Text.pack (ModuleName -> String
Plugins.moduleNameString ModuleName
moduleName_)

        (IO Context
getCurrentContext, IO ()
firstPluginPass, IO ()
lastPluginPass) <- do
            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
Shared.getTopLevelContext Text
moduleText)

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

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

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

                    Bool
isRoot <- String -> IO Bool
Shared.isRootModule (ModuleName -> String
Plugins.moduleNameString ModuleName
moduleName_)

                    -- Flush metrics if we're compiling one of the root
                    -- modules.  This is to work around the fact that we don't
                    -- have a proper way to finalize the `TracerProvider`
                    -- (since the finalizer would normally be responsible for
                    -- flushing any last metrics).
                    --
                    -- You might wonder: why don't we end the top-level span
                    -- here?  Well, we don't know which one of the root modules
                    -- will be the last one to be compiled.  However, flushing
                    -- once per root module is still fine because flushing is
                    -- safe to run at any time and in practice there will only
                    -- be a few root modules.
                    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
Monad.when Bool
isRoot IO ()
Shared.flush

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

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

        forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ CoreToDo
firstPass ] forall a. Semigroup a => a -> a -> a
<> [CoreToDo]
newTodos forall a. Semigroup a => a -> a -> a
<> [ CoreToDo
lastPass ])

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