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

{-| 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.
-}
module OpenTelemetry.Plugin.Shared
    ( -- * Plugin passes
      makeWrapperPluginPasses

      -- * Top-level context
    , initializeTopLevelContext
    , getTopLevelContext

      -- * Root module names
    , setRootModuleNames
    , isRootModule

      -- * Flushing
    , flush

    , getSampler
    ) where

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 OpenTelemetry.Trace
    ( Attribute(..)
    , PrimitiveAttribute(..)
    , InstrumentationLibrary(..)
    , Span
    , SpanArguments(..)
    , SpanContext(..)
    , Tracer
    , TracerProvider
    , TracerProviderOptions(..)
    )

import qualified Control.Concurrent.MVar as MVar
import qualified Data.HashMap.Strict as HashMap
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

{-| 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"

    forall (f :: * -> *) a. Applicative f => a -> f a
pure do
        String
"spanratio" <- Maybe String
maybeSampler
        String
ratioString <- Maybe String
maybeRatio
        Double
ratio <- forall a. Read a => String -> Maybe a
Read.readMaybe String
ratioString
        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 = forall a. IO a -> a
Unsafe.unsafePerformIO IO GenIO
MWC.createSystemRandom
{-# NOINLINE generator #-}

spanRatioBased :: Double -> Sampler
spanRatioBased :: Double -> Sampler
spanRatioBased Double
fraction = Sampler
    { getDescription :: Text
getDescription =
          Text
"SpanRatioBased{" forall a. Semigroup a => a -> a -> a
<> String -> Text
Text.pack (forall a. Show a => a -> String
show Double
fraction) 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 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 <- forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
(a, a) -> Gen (PrimState m) -> m a
MWC.uniformR (Double
0, Double
1) GenIO
generator

                let samplingResult :: SamplingResult
samplingResult =
                        if Double
random 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 ->
                        forall (f :: * -> *) a. Applicative f => a -> f a
pure TraceState
TraceState.empty

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

                forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingResult
samplingResult, 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 = forall a. IO a -> a
Unsafe.unsafePerformIO do
    ([Processor]
processors, TracerProviderOptions
options) <- 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
tracerProviderOptionsSampler = Sampler
sampler }

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

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

    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 ead 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 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
    MVar Span
spanMVar           <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a. IO (MVar a)
MVar.newEmptyMVar
    MVar Context
currentContextMVar <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO 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 Text Attribute
attributes =
                                forall k v. Hashable k => k -> v -> HashMap k v
HashMap.singleton Text
"sample" (PrimitiveAttribute -> Attribute
AttributeValue (Bool -> PrimitiveAttribute
BoolAttribute Bool
True))
                            }
                    else
                        SpanArguments
Trace.defaultSpanArguments

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

            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
_ <- forall a. MVar a -> a -> IO Bool
MVar.tryPutMVar MVar Context
currentContextMVar Context
currentContext

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

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

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

    forall (f :: * -> *) a. Applicative f => a -> f a
pure (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 = forall a. IO a -> a
Unsafe.unsafePerformIO forall a. IO (MVar a)
MVar.newEmptyMVar
{-# NOINLINE topLevelContextMVar #-}

getTopLevelSpan :: IO Span
getTopLevelSpan :: IO Span
getTopLevelSpan = 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 ->
            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 <- forall (m :: * -> *). MonadIO m => m Timestamp
Trace.Core.getTimestamp

            let arguments :: SpanArguments
arguments =
                    SpanArguments
Trace.defaultSpanArguments
                        { startTime :: Maybe Timestamp
startTime = forall a. a -> Maybe a
Just Timestamp
timestamp }

            Span
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

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

            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 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ByteString -> Maybe Baggage
W3CBaggage.decodeBaggage of
        Maybe Baggage
Nothing      -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Context
Context.empty
        Just Baggage
baggage -> 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 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (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 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 :: IO ()
initializeTopLevelContext :: IO ()
initializeTopLevelContext = do
    Span
span <- IO Span
getTopLevelSpan

    Context
context <- IO Context
getTopLevelBaggage

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

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

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

-- | Access the top-level `Context` computed by `initializeTopLevelContext`
getTopLevelContext :: IO Context
getTopLevelContext :: IO Context
getTopLevelContext = 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 = forall a. IO a -> a
Unsafe.unsafePerformIO 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 = forall a. Ord a => [a] -> Set a
Set.fromList (forall a b. (a -> b) -> [a] -> [b]
map String -> Text
Text.pack [String]
rootModuleNames)

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

    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 <- forall a. MVar a -> IO a
MVar.readMVar MVar (Set Text)
rootModuleNamesMVar

    forall (f :: * -> *) a. Applicative f => a -> f a
pure (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
_ <- forall (m :: * -> *).
MonadIO m =>
TracerProvider -> Maybe Int -> m FlushResult
Trace.Core.forceFlushTracerProvider TracerProvider
tracerProvider 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
_ <- forall (m :: * -> *).
MonadIO m =>
TracerProvider -> Maybe Int -> m FlushResult
Trace.Core.forceFlushTracerProvider TracerProvider
tracerProvider forall a. Maybe a
Nothing

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