{-# LANGUAGE BlockArguments    #-}
{-# LANGUAGE DerivingStrategies    #-}
{-# LANGUAGE GeneralizedNewtypeDeriving    #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

{-| This module provides the GHC-API-agnostic logic for this plugin (mostly
    open telemetry utilities)

    Because of how GHC plugins work, this module has to do some evil stuff
    under the hood to work within the confines of the plugin API.  That means
    that you should take care to use the utilities in this module correctly
    in order to avoid the plugin hanging.

    This module is intended to be shared by multiple GHC versions. Please
    don't depend on any GHC internal modules.
-}
module OpenTelemetry.Plugin.Shared
    ( -- * Plugin passes
      makeWrapperPluginPasses
    , getPluginShouldRecordPasses

      -- * Top-level context
    , initializeTopLevelContext
    , getTopLevelContext
    , modifyContextWithParentSpan
    , PackageName (..)

      -- * Root module names
    , setRootModuleNames
    , isRootModule

      -- * Flushing
    , flush
    , flushMetricsWhenRootModule

    , getSampler
    , tracer

    -- * Recording spans in 'runPhaseHook'
    , SpanMap
    , newSpanMap
    , recordModuleStart
    , recordModuleEnd
    ) where

import Control.Applicative ((<|>))
import Control.Concurrent.MVar (MVar)
import Control.Monad.IO.Class (liftIO)
import Data.ByteString (ByteString)
import Data.Set (Set)
import Data.Text (Text)
import OpenTelemetry.Context (Context)
import OpenTelemetry.Trace.Sampler (Sampler(..), SamplingResult(..))
import Prelude hiding (span)
import System.Random.MWC (GenIO)
import qualified StmContainers.Map as StmMap

import OpenTelemetry.Trace
    ( Attribute(..)
    , PrimitiveAttribute(..)
    , InstrumentationLibrary(..)
    , Span
    , SpanArguments(..)
    , SpanContext(..)
    , Tracer
    , TracerProvider
    , TracerProviderOptions(..)
    )

import qualified Control.Monad as Monad
import qualified Control.Concurrent.MVar as MVar
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Maybe as Maybe
import qualified Data.Set as Set
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text.Encoding
import qualified Data.Version as Version
import qualified OpenTelemetry.Context as Context
import qualified OpenTelemetry.Propagator.W3CBaggage as W3CBaggage
import qualified OpenTelemetry.Propagator.W3CTraceContext as W3CTraceContext
import qualified OpenTelemetry.Trace as Trace
import qualified OpenTelemetry.Trace.Core as Trace.Core
import qualified OpenTelemetry.Trace.Sampler as Sampler
import qualified OpenTelemetry.Trace.TraceState as TraceState
import qualified Paths_opentelemetry_plugin as Paths
import qualified System.Environment as Environment
import qualified System.IO.Unsafe as Unsafe
import qualified System.Random.MWC as MWC
import qualified Text.Read as Read
import qualified Control.Concurrent.STM as STM

{-| Very large Haskell builds can generate an enormous number of spans,
    but none of the stock samplers provide a way to sample a subset of
    the `Span`s within a trace.

    This adds a new "spanratio" sampler which can be used to sample subset of
    module spans.
-}
getSampler :: IO (Maybe Sampler)
getSampler :: IO (Maybe Sampler)
getSampler = do
    Maybe String
maybeSampler <- String -> IO (Maybe String)
Environment.lookupEnv String
"OTEL_TRACES_SAMPLER"

    Maybe String
maybeRatio <- String -> IO (Maybe String)
Environment.lookupEnv String
"OTEL_TRACES_SAMPLER_ARG"

    Maybe Sampler -> IO (Maybe Sampler)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure do
        String
"spanratio" <- Maybe String
maybeSampler
        String
ratioString <- Maybe String
maybeRatio
        Double
ratio <- String -> Maybe Double
forall a. Read a => String -> Maybe a
Read.readMaybe String
ratioString
        Sampler -> Maybe Sampler
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> Sampler
spanRatioBased Double
ratio)

{-| Like a lot of other uses of `Unsafe.unsafePerformIO` in this module, we're
    doing this because the plugin interface doesn't provide a way for us to
    acquire resources before returning the plugin.
-}
generator :: GenIO
generator :: GenIO
generator = IO (Gen RealWorld) -> Gen RealWorld
forall a. IO a -> a
Unsafe.unsafePerformIO IO (Gen RealWorld)
IO GenIO
MWC.createSystemRandom
{-# NOINLINE generator #-}

spanRatioBased :: Double -> Sampler
spanRatioBased :: Double -> Sampler
spanRatioBased Double
fraction = Sampler
    { getDescription :: Text
getDescription =
          Text
"SpanRatioBased{" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
Text.pack (Double -> String
forall a. Show a => a -> String
show Double
fraction) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"}"
    , shouldSample :: Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
shouldSample = \Context
context TraceId
traceId_ Text
name SpanArguments
spanArguments -> do
        case Text -> HashMap Text Attribute -> Maybe Attribute
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Text
"sample" (SpanArguments -> HashMap Text Attribute
attributes SpanArguments
spanArguments) of
            Just (AttributeValue (BoolAttribute Bool
True)) -> do
                Double
random <- (Double, Double) -> GenIO -> IO Double
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
(a, a) -> Gen (PrimState m) -> m a
forall (m :: * -> *).
PrimMonad m =>
(Double, Double) -> Gen (PrimState m) -> m Double
MWC.uniformR (Double
0, Double
1) GenIO
generator

                let samplingResult :: SamplingResult
samplingResult =
                        if Double
random Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
fraction then SamplingResult
RecordAndSample else SamplingResult
Drop

                TraceState
traceState_ <- case Context -> Maybe Span
Context.lookupSpan Context
context of
                    Maybe Span
Nothing ->
                        TraceState -> IO TraceState
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TraceState
TraceState.empty

                    Just Span
span ->
                        (SpanContext -> TraceState) -> IO SpanContext -> IO TraceState
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SpanContext -> TraceState
traceState (Span -> IO SpanContext
forall (m :: * -> *). MonadIO m => Span -> m SpanContext
Trace.Core.getSpanContext Span
span)

                (SamplingResult, HashMap Text Attribute, TraceState)
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingResult
samplingResult, HashMap Text Attribute
forall k v. HashMap k v
HashMap.empty, TraceState
traceState_)

            Maybe Attribute
_ ->
                Sampler
-> Context
-> TraceId
-> Text
-> SpanArguments
-> IO (SamplingResult, HashMap Text Attribute, TraceState)
shouldSample Sampler
Sampler.alwaysOn Context
context TraceId
traceId_ Text
name SpanArguments
spanArguments
    }

{-| Note: We don't properly shut this down using `Trace.shutdownTracerProvider`,
    but all that the shutdown does is flush metrics, so instead we flush metrics
    (using `flush`) at the end of compilation to make up for the lack of a
    proper shutdown.
-}
tracerProvider :: TracerProvider
tracerProvider :: TracerProvider
tracerProvider = IO TracerProvider -> TracerProvider
forall a. IO a -> a
Unsafe.unsafePerformIO do
    ([Processor]
processors, TracerProviderOptions
options) <-
        -- This function will collect *all* of the command line arguments
        -- that were provided to GHC. This results in a huge amount of data
        -- being sent. For that reason, we blank out the process arguments
        -- for this section of code.
        [String]
-> IO ([Processor], TracerProviderOptions)
-> IO ([Processor], TracerProviderOptions)
forall a. [String] -> IO a -> IO a
Environment.withArgs [] do
            IO ([Processor], TracerProviderOptions)
Trace.getTracerProviderInitializationOptions

    Maybe Sampler
maybeSampler <- IO (Maybe Sampler)
getSampler

    let newOptions :: TracerProviderOptions
newOptions =
            case Maybe Sampler
maybeSampler of
                Maybe Sampler
Nothing      -> TracerProviderOptions
options
                Just Sampler
sampler -> TracerProviderOptions
options{ tracerProviderOptionsSampler = sampler }

    TracerProvider
tracerProvider_ <- [Processor] -> TracerProviderOptions -> IO TracerProvider
forall (m :: * -> *).
MonadIO m =>
[Processor] -> TracerProviderOptions -> m TracerProvider
Trace.createTracerProvider [Processor]
processors TracerProviderOptions
newOptions

    TracerProvider -> IO ()
forall (m :: * -> *). MonadIO m => TracerProvider -> m ()
Trace.setGlobalTracerProvider TracerProvider
tracerProvider_

    TracerProvider -> IO TracerProvider
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TracerProvider
tracerProvider_
{-# NOINLINE tracerProvider #-}

tracer :: Tracer
tracer :: Tracer
tracer =
    TracerProvider -> InstrumentationLibrary -> TracerOptions -> Tracer
Trace.makeTracer TracerProvider
tracerProvider InstrumentationLibrary
instrumentationLibrary TracerOptions
Trace.tracerOptions
  where
    instrumentationLibrary :: InstrumentationLibrary
instrumentationLibrary =
        InstrumentationLibrary
            { libraryName :: Text
libraryName    = Text
"opentelemetry-plugin"
            , libraryVersion :: Text
libraryVersion = String -> Text
Text.pack (Version -> String
Version.showVersion Version
Paths.version)
            }

{-| This used by the GHC plugin to create two plugin passes that start and stop
    a `Span`, respectively.

    In order for `Span` ancestry to be tracked correctly this takes an
    @`IO` `Context`@ as an input (to read the parent `Span`'s `Context`) and
    returns an @`IO` `Context`@ as an output (to read the current `Span`'s
    `Context`).
-}
makeWrapperPluginPasses
    :: Bool
      -- ^ Whether to sample a subset of spans
    -> IO Context
       -- ^ Action to read the parent span's `Context`
    -> Text
       -- ^ Label for the current span
    -> IO (IO Context, IO (), IO ())
makeWrapperPluginPasses :: Bool -> IO Context -> Text -> IO (IO Context, IO (), IO ())
makeWrapperPluginPasses Bool
sample IO Context
getParentContext Text
label = IO (IO Context, IO (), IO ()) -> IO (IO Context, IO (), IO ())
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
    MVar Span
spanMVar           <- IO (MVar Span) -> IO (MVar Span)
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (MVar Span)
forall a. IO (MVar a)
MVar.newEmptyMVar
    MVar Context
currentContextMVar <- IO (MVar Context) -> IO (MVar Context)
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (MVar Context)
forall a. IO (MVar a)
MVar.newEmptyMVar

    let beginPass :: IO ()
beginPass = do
            Context
parentContext <- IO Context
getParentContext

            let spanArguments :: SpanArguments
spanArguments =
                    if Bool
sample
                    then
                        SpanArguments
Trace.defaultSpanArguments
                            { attributes =
                                HashMap.singleton "sample" (AttributeValue (BoolAttribute True))
                            }
                    else
                        SpanArguments
Trace.defaultSpanArguments

            Span
passSpan <- Tracer -> Context -> Text -> SpanArguments -> IO Span
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
Tracer -> Context -> Text -> SpanArguments -> m Span
Trace.createSpan Tracer
tracer Context
parentContext Text
label SpanArguments
spanArguments

            Bool
_ <- MVar Span -> Span -> IO Bool
forall a. MVar a -> a -> IO Bool
MVar.tryPutMVar MVar Span
spanMVar Span
passSpan

            let currentContext :: Context
currentContext = Span -> Context -> Context
Context.insertSpan Span
passSpan Context
parentContext

            Bool
_ <- MVar Context -> Context -> IO Bool
forall a. MVar a -> a -> IO Bool
MVar.tryPutMVar MVar Context
currentContextMVar Context
currentContext

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

    let endPass :: IO ()
endPass = do
            Span
passSpan <- MVar Span -> IO Span
forall a. MVar a -> IO a
MVar.readMVar MVar Span
spanMVar

            Span -> Maybe Timestamp -> IO ()
forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
Trace.endSpan Span
passSpan Maybe Timestamp
forall a. Maybe a
Nothing

    (IO Context, IO (), IO ()) -> IO (IO Context, IO (), IO ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MVar Context -> IO Context
forall a. MVar a -> IO a
MVar.readMVar MVar Context
currentContextMVar, IO ()
beginPass, IO ()
endPass)

{-| We're intentionally **NOT** using `OpenTelemetry.Context.ThreadLocal`
    here since the `GHC.Plugins.Plugin` logic doesn't necessarily run in a
    single thread (@ghc@ builds can be multi-threaded).  Instead, we provide
    our own `Context` global variable.
-}
topLevelContextMVar :: MVar Context
topLevelContextMVar :: MVar Context
topLevelContextMVar = IO (MVar Context) -> MVar Context
forall a. IO a -> a
Unsafe.unsafePerformIO IO (MVar Context)
forall a. IO (MVar a)
MVar.newEmptyMVar
{-# NOINLINE topLevelContextMVar #-}

{- | We keep track of the module spans in this top-level 'MVar' so that it
     may be shared between the driverPlugin and other plugins.

-}
topLevelSpanMapMVar :: MVar SpanMap
topLevelSpanMapMVar :: MVar SpanMap
topLevelSpanMapMVar = IO (MVar SpanMap) -> MVar SpanMap
forall a. IO a -> a
Unsafe.unsafePerformIO IO (MVar SpanMap)
forall a. IO (MVar a)
MVar.newEmptyMVar
{-# NOINLINE topLevelSpanMapMVar #-}

getTopLevelSpan
    :: PackageName
    -> IO Span
getTopLevelSpan :: PackageName -> IO Span
getTopLevelSpan PackageName
packageName = do
    Maybe ByteString
traceParent <- String -> IO (Maybe ByteString)
lookupEnv String
"TRACEPARENT"
    Maybe ByteString
traceState_ <- String -> IO (Maybe ByteString)
lookupEnv String
"TRACESTATE"

    case Maybe ByteString -> Maybe ByteString -> Maybe SpanContext
W3CTraceContext.decodeSpanContext Maybe ByteString
traceParent Maybe ByteString
traceState_ of
        Just SpanContext
spanContext ->
            Span -> IO Span
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SpanContext -> Span
Trace.Core.wrapSpanContext SpanContext
spanContext)

        Maybe SpanContext
Nothing -> do
            -- If we're not inheriting a span from
            -- `TRACEPARENT`/`TRACESTATE`, then create a zero-duration span
            -- whose sole purpose is to be a parent span for each module's
            -- spans.
            --
            -- Ideally we'd like this span's duration to last for the
            -- entirety of compilation, but there isn't a good way to end
            -- the span when compilation is done.  Also, we still need
            -- *some* parent span for each module's spans, otherwise an
            -- entirely new trace will be created for each new span.
            -- Creating a zero-duration span is the least-worst solution.
            --
            -- Note that there aren't any issues with the child spans
            -- lasting longer than the parent span.  This is supported by
            -- open telemetry and the Haskell API.
            Timestamp
timestamp <- IO Timestamp
forall (m :: * -> *). MonadIO m => m Timestamp
Trace.Core.getTimestamp

            let arguments :: SpanArguments
arguments =
                    SpanArguments
Trace.defaultSpanArguments
                        { startTime = Just timestamp
                        , attributes = HashMap.fromList
                            [ ("packageName", Trace.Core.toAttribute packageName)
                            ]
                        }

            Span
span <- Tracer -> Context -> Text -> SpanArguments -> IO Span
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
Tracer -> Context -> Text -> SpanArguments -> m Span
Trace.createSpan Tracer
tracer Context
Context.empty Text
"opentelemetry GHC plugin" SpanArguments
arguments

            Span -> Maybe Timestamp -> IO ()
forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
Trace.endSpan Span
span (Timestamp -> Maybe Timestamp
forall a. a -> Maybe a
Just Timestamp
timestamp)

            Span -> IO Span
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Span
span

getTopLevelBaggage :: IO Context
getTopLevelBaggage :: IO Context
getTopLevelBaggage = do
    Maybe ByteString
maybeBytes <- String -> IO (Maybe ByteString)
lookupEnv String
"BAGGAGE"
    case Maybe ByteString
maybeBytes Maybe ByteString -> (ByteString -> Maybe Baggage) -> Maybe Baggage
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ByteString -> Maybe Baggage
W3CBaggage.decodeBaggage of
        Maybe Baggage
Nothing      -> Context -> IO Context
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Context
Context.empty
        Just Baggage
baggage -> Context -> IO Context
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Baggage -> Context -> Context
Context.insertBaggage Baggage
baggage Context
Context.empty)

lookupEnv :: String -> IO (Maybe ByteString)
lookupEnv :: String -> IO (Maybe ByteString)
lookupEnv = (IO (Maybe String) -> IO (Maybe ByteString))
-> (String -> IO (Maybe String)) -> String -> IO (Maybe ByteString)
forall a b. (a -> b) -> (String -> a) -> String -> b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Maybe String -> Maybe ByteString)
-> IO (Maybe String) -> IO (Maybe ByteString)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((String -> ByteString) -> Maybe String -> Maybe ByteString
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap String -> ByteString
encode)) String -> IO (Maybe String)
Environment.lookupEnv
  where
    encode :: String -> ByteString
encode = Text -> ByteString
Text.Encoding.encodeUtf8 (Text -> ByteString) -> (String -> Text) -> String -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
Text.pack

{-| This initializes the top-level `Context` using the @TRACEPARENT@ \/
    @TRACESTATE@ \/ @BAGGAGE@ environment variables (if present) and otherwise
    sets it to the empty `Context`

    You have to run this command before calling `getTopLevelContext` otherwise
    the latter will hang.
-}
initializeTopLevelContext
    :: PackageName
    -> IO ()
initializeTopLevelContext :: PackageName -> IO ()
initializeTopLevelContext PackageName
packageName = do
    Span
span <- PackageName -> IO Span
getTopLevelSpan PackageName
packageName

    Context
context <- IO Context
getTopLevelBaggage

    let contextWithSpan :: Context
contextWithSpan = Span -> Context -> Context
Context.insertSpan Span
span Context
context

    Bool
_ <- MVar Context -> Context -> IO Bool
forall a. MVar a -> a -> IO Bool
MVar.tryPutMVar MVar Context
topLevelContextMVar Context
contextWithSpan

    Bool
_ <- MVar SpanMap -> SpanMap -> IO Bool
forall a. MVar a -> a -> IO Bool
MVar.tryPutMVar MVar SpanMap
topLevelSpanMapMVar (SpanMap -> IO Bool) -> IO SpanMap -> IO Bool
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO SpanMap
newSpanMap

    () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Access the top-level `Context` computed by `initializeTopLevelContext`
getTopLevelContext :: IO Context
getTopLevelContext :: IO Context
getTopLevelContext = MVar Context -> IO Context
forall a. MVar a -> IO a
MVar.readMVar MVar Context
topLevelContextMVar

{-| This is used for communicating between `GHC.Plugins.driverPlugin` and
    `GHC.Plugins.installCoreToDos`, because only `GHC.Plugins.driverPlugin` has
    access to the full module graph, but there isn't a good way within the
    `GHC.Plugins.Plugin` API to share that information with the rest of the
    plugin other than a global variable.
-}
rootModuleNamesMVar :: MVar (Set Text)
rootModuleNamesMVar :: MVar (Set Text)
rootModuleNamesMVar = IO (MVar (Set Text)) -> MVar (Set Text)
forall a. IO a -> a
Unsafe.unsafePerformIO IO (MVar (Set Text))
forall a. IO (MVar a)
MVar.newEmptyMVar
{-# NOINLINE rootModuleNamesMVar #-}

{-| Set the root module names (computed by `GHC.Plugins.driverPlugin`)

    You have to run this command before calling `isRootModule` otherwise
    the latter will hang.
-}
setRootModuleNames :: [String] -> IO ()
setRootModuleNames :: [String] -> IO ()
setRootModuleNames [String]
rootModuleNames = do
    let set :: Set Text
set = [Text] -> Set Text
forall a. Ord a => [a] -> Set a
Set.fromList ((String -> Text) -> [String] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map String -> Text
Text.pack [String]
rootModuleNames)

    Bool
_ <- MVar (Set Text) -> Set Text -> IO Bool
forall a. MVar a -> a -> IO Bool
MVar.tryPutMVar MVar (Set Text)
rootModuleNamesMVar Set Text
set

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

-- | Check if a module is one of the root modules
isRootModule :: String -> IO Bool
isRootModule :: String -> IO Bool
isRootModule String
moduleName = do
    Set Text
rootModuleNames <- MVar (Set Text) -> IO (Set Text)
forall a. MVar a -> IO a
MVar.readMVar MVar (Set Text)
rootModuleNamesMVar

    Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> Set Text -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member (String -> Text
Text.pack String
moduleName) Set Text
rootModuleNames)

-- | Flush all metrics
flush :: IO ()
flush :: IO ()
flush = do
    FlushResult
_ <- TracerProvider -> Maybe Int -> IO FlushResult
forall (m :: * -> *).
MonadIO m =>
TracerProvider -> Maybe Int -> m FlushResult
Trace.Core.forceFlushTracerProvider TracerProvider
tracerProvider Maybe Int
forall a. Maybe a
Nothing
    -- We can't check the result yet because
    -- `FlushResult` is not exported by
    -- `hs-opentelemetry-api`
    --
    -- https://github.com/iand675/hs-opentelemetry/pull/96
    FlushResult
_ <- TracerProvider -> Maybe Int -> IO FlushResult
forall (m :: * -> *).
MonadIO m =>
TracerProvider -> Maybe Int -> m FlushResult
Trace.Core.forceFlushTracerProvider TracerProvider
tracerProvider Maybe Int
forall a. Maybe a
Nothing

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

-- | Returns 'True' if the plugin should create spans for module passes in
-- compilation. Examples would be Simplifier, any other plugin execution,
-- etc.
getPluginShouldRecordPasses :: IO Bool
getPluginShouldRecordPasses :: IO Bool
getPluginShouldRecordPasses = do
    Maybe String
maybeRecordPasses <- String -> IO (Maybe String)
Environment.lookupEnv String
"OTEL_GHC_PLUGIN_RECORD_PASSES"
    Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
Maybe.fromMaybe Bool
False do
        String
recordPasses <- Maybe String
maybeRecordPasses
        Bool -> Maybe Bool
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> Maybe Bool) -> Bool -> Maybe Bool
forall a b. (a -> b) -> a -> b
$ Text -> Text
Text.toLower (String -> Text
Text.pack String
recordPasses) Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text
"true", Text
"t"]


-- | 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.
flushMetricsWhenRootModule :: String -> IO ()
flushMetricsWhenRootModule :: String -> IO ()
flushMetricsWhenRootModule String
modName = do
    Bool
isRoot <- String -> IO Bool
isRootModule String
modName
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
Monad.when Bool
isRoot IO ()
flush

-- | A concurrently accessible map that can be used to connect a module at
-- beginning of compilation and at the end.
--
-- GHC records the phases of computation in a datatype 'TPhase'. This
-- datatype begins Haskell compilation with the 'T_Hsc' phase. The final
-- phase in compilation is 'T_MergeForeign'. The final phase has a few
-- items: a 'PipeEnv', an 'HscEnv', a 'Filepath' representing the location
-- of the object file for the module, and a list of 'Filepath' that I don't
-- know the purpose of.
--
-- 'T_Hsc' phase carries a 'ModSummary' type, which fortunately includes
-- a 'ms_location' field which has the 'ml_object_file' field. Since this
-- information is present both at the beginning and end, we can use that to
-- associate a 'Trace.Span' with a module's beginning and end, to record
-- the full time in compilation.
data SpanMap = SpanMap
    { SpanMap -> Map String SpanRelease
fromObjectFile :: StmMap.Map FilePath SpanRelease
    , SpanMap -> Map String SpanRelease
fromModuleName :: StmMap.Map String SpanRelease
    }

-- | A 'Trace.Span' coupled with the action to delete the 'Span' from the
-- 'SpanMap'.
data SpanRelease = SpanRelease
    { SpanRelease -> Span
spanReleaseSpan :: Trace.Span
    , SpanRelease -> STM ()
spanReleaseAction :: STM.STM ()
    , SpanRelease -> String
spanReleaseModuleName :: String
    }

-- | Create a new empty 'SpanMap'.
newSpanMap :: IO SpanMap
newSpanMap :: IO SpanMap
newSpanMap = Map String SpanRelease -> Map String SpanRelease -> SpanMap
SpanMap (Map String SpanRelease -> Map String SpanRelease -> SpanMap)
-> IO (Map String SpanRelease)
-> IO (Map String SpanRelease -> SpanMap)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Map String SpanRelease)
forall key value. IO (Map key value)
StmMap.newIO IO (Map String SpanRelease -> SpanMap)
-> IO (Map String SpanRelease) -> IO SpanMap
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IO (Map String SpanRelease)
forall key value. IO (Map key value)
StmMap.newIO

-- | Create a 'Span' for the given 'ModSummary' and record it in the
-- 'SpanMap'.
recordModuleStart
    :: PackageName
    -> FilePath
    -- ^ The location of the object file for the module. This should be
    -- available through the 'ModSummary' via 'ms_location' and
    -- 'ml_obj_file'
    -> String
    -- ^ A string representing the name of the module.
    -> IO ()
recordModuleStart :: PackageName -> String -> String -> IO ()
recordModuleStart PackageName
packageName String
modObjectLocation String
modName = do
    SpanMap
spanMap <- MVar SpanMap -> IO SpanMap
forall a. MVar a -> IO a
MVar.readMVar MVar SpanMap
topLevelSpanMapMVar
    Context
context <- IO Context
getTopLevelContext
    Span
span_ <- Tracer -> Context -> Text -> SpanArguments -> IO Span
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
Tracer -> Context -> Text -> SpanArguments -> m Span
Trace.createSpan Tracer
tracer Context
context (String -> Text
Text.pack String
modName) SpanArguments
Trace.defaultSpanArguments
        { attributes = HashMap.fromList
            [ ("packageName", Trace.Core.toAttribute packageName)
            ]
        }
    let spanRelease :: SpanRelease
spanRelease =
            SpanRelease
                { spanReleaseSpan :: Span
spanReleaseSpan =
                    Span
span_
                , spanReleaseAction :: STM ()
spanReleaseAction = do
                    String -> Map String SpanRelease -> STM ()
forall key value. Hashable key => key -> Map key value -> STM ()
StmMap.delete String
modObjectLocation (SpanMap -> Map String SpanRelease
fromObjectFile SpanMap
spanMap)
                    String -> Map String SpanRelease -> STM ()
forall key value. Hashable key => key -> Map key value -> STM ()
StmMap.delete String
modName (SpanMap -> Map String SpanRelease
fromModuleName SpanMap
spanMap)
                , spanReleaseModuleName :: String
spanReleaseModuleName =
                    String
modName
                }
    STM () -> IO ()
forall a. STM a -> IO a
STM.atomically do
        SpanRelease -> String -> Map String SpanRelease -> STM ()
forall key value.
Hashable key =>
value -> key -> Map key value -> STM ()
StmMap.insert SpanRelease
spanRelease String
modObjectLocation (SpanMap -> Map String SpanRelease
fromObjectFile SpanMap
spanMap)
        SpanRelease -> String -> Map String SpanRelease -> STM ()
forall key value.
Hashable key =>
value -> key -> Map key value -> STM ()
StmMap.insert SpanRelease
spanRelease String
modName (SpanMap -> Map String SpanRelease
fromModuleName SpanMap
spanMap)

-- | Given a 'Plugins.Module' and a function that provides
-- a 'Trace.Context', this function modifies the 'Trace.Context' to have
-- the parent span of the module.
modifyContextWithParentSpan
    :: String
    -> IO Context.Context
    -> IO Context.Context
modifyContextWithParentSpan :: String -> IO Context -> IO Context
modifyContextWithParentSpan String
module_ IO Context
getContext = do
    Maybe Span
mspan <- String -> IO (Maybe Span)
getSpanForModule String
module_
    let insertSpan :: Context -> Context
insertSpan Context
context =
            Context -> (Span -> Context) -> Maybe Span -> Context
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Context
context (Span -> Context -> Context
`Context.insertSpan` Context
context) Maybe Span
mspan
    (Context -> Context) -> IO Context -> IO Context
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Context -> Context
insertSpan IO Context
getContext

-- | Retrieve the 'Trace.Span' for a given 'Plugins.Module', if one has
-- been recorded.
getSpanForModule
    :: String
    -> IO (Maybe Span)
getSpanForModule :: String -> IO (Maybe Span)
getSpanForModule String
moduleNameString = do
    SpanMap
spanMap <- MVar SpanMap -> IO SpanMap
forall a. MVar a -> IO a
MVar.readMVar MVar SpanMap
topLevelSpanMapMVar
    STM (Maybe Span) -> IO (Maybe Span)
forall a. STM a -> IO a
STM.atomically do
        (SpanRelease -> Span) -> Maybe SpanRelease -> Maybe Span
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SpanRelease -> Span
spanReleaseSpan (Maybe SpanRelease -> Maybe Span)
-> STM (Maybe SpanRelease) -> STM (Maybe Span)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> do
            String -> Map String SpanRelease -> STM (Maybe SpanRelease)
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
StmMap.lookup String
moduleNameString (Map String SpanRelease -> STM (Maybe SpanRelease))
-> Map String SpanRelease -> STM (Maybe SpanRelease)
forall a b. (a -> b) -> a -> b
$ SpanMap -> Map String SpanRelease
fromModuleName SpanMap
spanMap

recordModuleEnd
    :: String
    -- ^ This should be either the module name *or* the object file
    -- location.
    ->  IO ()
recordModuleEnd :: String -> IO ()
recordModuleEnd String
moduleIdentifier = do
    SpanMap
spanMap <- MVar SpanMap -> IO SpanMap
forall a. MVar a -> IO a
MVar.readMVar MVar SpanMap
topLevelSpanMapMVar
    Maybe SpanRelease
mspan <- STM (Maybe SpanRelease) -> IO (Maybe SpanRelease)
forall a. STM a -> IO a
STM.atomically do
        (Maybe SpanRelease -> Maybe SpanRelease -> Maybe SpanRelease)
-> STM (Maybe SpanRelease)
-> STM (Maybe SpanRelease)
-> STM (Maybe SpanRelease)
forall a b c. (a -> b -> c) -> STM a -> STM b -> STM c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2
            Maybe SpanRelease -> Maybe SpanRelease -> Maybe SpanRelease
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
(<|>)
            do String -> Map String SpanRelease -> STM (Maybe SpanRelease)
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
StmMap.lookup String
moduleIdentifier (SpanMap -> Map String SpanRelease
fromObjectFile SpanMap
spanMap)
            do String -> Map String SpanRelease -> STM (Maybe SpanRelease)
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
StmMap.lookup String
moduleIdentifier (SpanMap -> Map String SpanRelease
fromModuleName SpanMap
spanMap)

    Maybe SpanRelease -> (SpanRelease -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
Monad.forM_ Maybe SpanRelease
mspan \SpanRelease {String
STM ()
Span
spanReleaseSpan :: SpanRelease -> Span
spanReleaseAction :: SpanRelease -> STM ()
spanReleaseModuleName :: SpanRelease -> String
spanReleaseSpan :: Span
spanReleaseAction :: STM ()
spanReleaseModuleName :: String
..} -> do
        Span -> Maybe Timestamp -> IO ()
forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
Trace.endSpan Span
spanReleaseSpan Maybe Timestamp
forall a. Maybe a
Nothing
        STM () -> IO ()
forall a. STM a -> IO a
STM.atomically STM ()
spanReleaseAction
        String -> IO ()
flushMetricsWhenRootModule String
spanReleaseModuleName

-- | The name of the package that is being compiled.
newtype PackageName = PackageName
    { PackageName -> Text
unPackageName :: Text
    }
    deriving newtype PackageName -> Attribute
(PackageName -> Attribute) -> ToAttribute PackageName
forall a. (a -> Attribute) -> ToAttribute a
$ctoAttribute :: PackageName -> Attribute
toAttribute :: PackageName -> Attribute
Trace.Core.ToAttribute