Effect.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE OverloadedLists #-}

module Effectful.OpenTelemetry.Tracing.Effect
    ( -- * Effect
      Tracing
    , inSpan
    , inOkSpan
    , inErrorSpan
    , withContext
    , addEvent
    , addEventNow
    , recordException
    , recordExceptionAt
    , currentContext
    , currentContextIO

      -- * Runners
    , runTracing
    , runHttpTracing
    , runGrpcTracing
    , runConsoleTracing
    , runInMemoryTracing
    , runTracingWith
    , runNoTracing
    )
where

import Control.Monad.Extra (ifM)
import Data.Aeson (ToJSON (..))
import Data.Functor (void, (<&>))
import Data.Sequence (Seq, (|>))
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Typeable (typeOf)
import Effectful
import Effectful.Concurrent (Concurrent)
import Effectful.Dispatch.Static
import Effectful.Environment (Environment)
import Effectful.Error.Static (Error, throwError, tryError)
import Effectful.Exception
    ( Exception
    , SomeException
    , displayException
    , mask
    , throwIO
    , try
    , uninterruptibleMask_
    )
import Effectful.Http2Client (HostName, PortNumber)
import Effectful.HttpClient (responseTimeoutDefault)
import Effectful.OpenTelemetry.Exporter (Exporter)
import Effectful.OpenTelemetry.Exporter qualified as Exporter
import Effectful.OpenTelemetry.Exporter.Console qualified as Console
import Effectful.OpenTelemetry.Exporter.Environment qualified as Exporter.Environment
import Effectful.OpenTelemetry.Exporter.OTLP (defaultGrpcSendTimeout)
import Effectful.OpenTelemetry.Protocol
    ( Attributes
    , Compression
    , Encoding
    , OTLP
    , Resource
    , Scope
    , exportIO
    , runInMemoryOTLP
    , runNoOTLP
    , runOTLPWith
    )
import Effectful.OpenTelemetry.Protocol.Environment qualified as Environment
import Effectful.OpenTelemetry.Protocol.Export qualified as Export
import Effectful.OpenTelemetry.Timestamp (Timestamp)
import Effectful.OpenTelemetry.Timestamp qualified as Timestamp
import Effectful.OpenTelemetry.Tracing.Span (Span (..))
import Effectful.OpenTelemetry.Tracing.Span.Context qualified as Span (Context)
import Effectful.OpenTelemetry.Tracing.Span.Context qualified as Span.Context
import Effectful.OpenTelemetry.Tracing.Span.Event (Event (..))
import Effectful.OpenTelemetry.Tracing.Span.Event qualified as Event
import Effectful.OpenTelemetry.Tracing.Span.Kind qualified as Span
import Effectful.OpenTelemetry.Tracing.Span.Status qualified as Span.Status
import Effectful.Retry (Retry)
import Effectful.Timeout (Timeout)
import GHC.Stack (whoCreated)
import Network.URI (URI)
import Prelude

data Tracing :: Effect

type instance DispatchOf Tracing = 'Static 'WithSideEffects

data instance StaticRep Tracing = Tracing
    { currentSpan :: Maybe Span.Context
    , events :: Seq Event
    , export :: Span -> IO ()
    }

data ExitCase e
    = ExitSuccess
    | ExitFailure e

-- | Execute an action with a given 'Span.Context'.
withContext :: (Tracing :> es) => Span.Context -> Eff es a -> Eff es a
withContext context = localStaticRep \tracing ->
    tracing{currentSpan = Just context, events = mempty}

-- | Execute an action within a named 'Span'.
-- If the action throws an exception, the 'Span' is finalised with 'Span.Status.Error'
-- and the exception is re-raised.
-- Otherwise, the 'Span' is exported with 'Span.Status.Unset'.
inSpan
    :: forall es a
     . (HasCallStack, Tracing :> es)
    => Text
    -> Span.Kind
    -> Attributes
    -> Eff es a
    -> Eff es a
inSpan name kind attributes action = do
    tracing <- getStaticRep @Tracing
    context <- unsafeEff_ $ Span.Context.new tracing.currentSpan
    startTime <- unsafeEff_ Timestamp.now
    withContext context $ finallyCaseIO @SomeException action \exitCase -> do
        endTime <- unsafeEff_ Timestamp.now
        tracing' <- getStaticRep @Tracing
        unsafeEff_ . tracing.export $
            Span
                { name
                , parentSpanId = tracing.currentSpan <&> (.spanId)
                , context
                , kind
                , startTime
                , endTime = Just endTime
                , attributes
                , events = tracing'.events
                , status =
                    case exitCase of
                        ExitSuccess -> Span.Status.unset
                        (ExitFailure e) -> Span.Status.fromException e
                }
  where
    finallyCaseIO
        :: forall e
         . (Exception e)
        => Eff es a
        -> (ExitCase e -> Eff es ())
        -> Eff es a
    finallyCaseIO thing after = mask \restore -> do
        try @e (restore thing) >>= \case
            Left e -> do
                recordException e mempty
                void . try @e . uninterruptibleMask_ . after $ ExitFailure e
                throwIO e
            Right y -> do
                void . try @e . uninterruptibleMask_ $ after ExitSuccess
                pure y

-- | Execute an action that may not fail within a named 'Span'.
-- Explicity marks the 'Span' with 'Span.Status.Ok'.
-- Does not export the 'Span' if the action raises an exception.
-- Use 'inErrorSpan' or 'inSpan' for actions that may fail.
inOkSpan
    :: forall a es
     . (Tracing :> es)
    => Text
    -> Span.Kind
    -> Attributes
    -> Eff es a
    -> Eff es a
inOkSpan name kind attributes action = do
    tracing <- getStaticRep @Tracing
    context <- unsafeEff_ $ Span.Context.new tracing.currentSpan
    startTime <- unsafeEff_ Timestamp.now
    (a, tracing') <- withContext context $ (,) <$> action <*> getStaticRep @Tracing
    endTime <- unsafeEff_ Timestamp.now
    unsafeEff_ . tracing.export $
        Span
            { name
            , parentSpanId = tracing.currentSpan <&> (.spanId)
            , context
            , kind
            , startTime
            , endTime = Just endTime
            , attributes
            , events = tracing'.events
            , status = Span.Status.ok
            }
    pure a

-- | Execute an action that may fail with the expected @Error e@ within a named 'Span'.
-- If the action fails with this expected @Error e@, the 'Span' is finalised with 'Span.Status.Error'
-- and the exception is re-raised.
-- Otherwise, exports the 'Span' with 'Span.Status.Unset'.
-- Does not export the 'Span' if the action fails with an unexpected exception (e.g. via 'throwIO').
-- If you need to export 'Span's in the presence of any exception, use 'inSpan' instead.
inErrorSpan
    :: forall e es a
     . (Exception e, Error e :> es, Tracing :> es)
    => Text
    -> Span.Kind
    -> Attributes
    -> Eff es a
    -> Eff es a
inErrorSpan name kind attributes action = do
    tracing <- getStaticRep @Tracing
    context <- unsafeEff_ $ Span.Context.new tracing.currentSpan
    startTime <- unsafeEff_ Timestamp.now
    withContext context . finallyCase action $ \exitCase -> do
        endTime <- unsafeEff_ Timestamp.now
        tracing' <- getStaticRep @Tracing
        unsafeEff_ . tracing.export $
            Span
                { name
                , parentSpanId = tracing.currentSpan <&> (.spanId)
                , context
                , kind
                , startTime
                , endTime = Just endTime
                , attributes
                , events = tracing'.events
                , status =
                    case exitCase of
                        ExitSuccess -> Span.Status.unset
                        (ExitFailure e) -> Span.Status.fromException e
                }
  where
    finallyCase :: Eff es a -> (ExitCase e -> Eff es ()) -> Eff es a
    finallyCase thing after = do
        tryError @e thing >>= \case
            Left (_, e) -> do
                recordException e mempty
                void . tryError @e . uninterruptibleMask_ . after $ ExitFailure e
                throwError e
            Right y -> do
                void . tryError @e . uninterruptibleMask_ . after $ ExitSuccess
                pure y

-- | Add an event to the current span. This is a no-op if not called in a span.
--
-- https://github.com/open-telemetry/opentelemetry-specification/blob/v1.57.0/specification/trace/api.md#add-events
addEvent :: (Tracing :> es) => Event -> Eff es ()
addEvent event = do
    tracing <- getStaticRep
    case tracing.currentSpan of
        Nothing -> pure ()
        Just context ->
            putStaticRep $ tracing{currentSpan = Just context, events = tracing.events |> event}

-- | Add an event that occurred at the current time to the current span. This is a no-op if not called in a span.
--
-- https://github.com/open-telemetry/opentelemetry-specification/blob/v1.57.0/specification/trace/api.md#add-events
addEventNow :: (Tracing :> es) => Text -> Attributes -> Eff es ()
addEventNow name attributes = addEvent =<< unsafeEff_ (Event.now name attributes)

-- | Add an exception to the current span. This is a no-op if not called in a span.
--
-- - https://github.com/open-telemetry/opentelemetry-specification/blob/v1.57.0/specification/trace/api.md#record-exception
-- - https://github.com/open-telemetry/opentelemetry-specification/blob/v1.57.0/specification/trace/exceptions.md
recordExceptionAt :: (Exception e, Tracing :> es) => Timestamp -> e -> Attributes -> Eff es ()
recordExceptionAt time e attributes = do
    stackTrace <- unsafeEff_ $ whoCreated e
    addEvent
        Event
            { time
            , name = "exception"
            , attributes =
                [ ("exception.type", toJSON . Text.pack . show $ typeOf e)
                , ("exception.message", toJSON . Text.pack $ displayException e)
                , ("exception.stacktrace", toJSON . Text.unlines $ Text.pack <$> stackTrace)
                ]
                    <> attributes
            }

-- | Add an exception event that occurred at the current time to the current span. This is a no-op if not called in a span.
--
-- - https://github.com/open-telemetry/opentelemetry-specification/blob/v1.57.0/specification/trace/api.md#record-exception
-- - https://github.com/open-telemetry/opentelemetry-specification/blob/v1.57.0/specification/trace/exceptions.md
recordException :: (Exception e, Tracing :> es) => e -> Attributes -> Eff es ()
recordException e attributes = do
    time <- unsafeEff_ Timestamp.now
    recordExceptionAt time e attributes

currentContext :: (Tracing :> es) => Eff es (Maybe Span.Context)
currentContext = do
    Tracing{..} <- getStaticRep
    pure currentSpan

currentContextIO :: (Tracing :> es) => Eff es (IO (Maybe Span.Context))
currentContextIO = unsafeEff $ pure . unEff currentContext

runTracingState
    :: (OTLP Span :> es, IOE :> es)
    => Eff (Tracing ': es) a
    -> Eff es a
runTracingState eff = do
    export <- exportIO
    evalStaticRep Tracing{currentSpan = Nothing, events = mempty, ..} eff

-- | Run the 'Tracing' effect, sending telemetry to an exporter.
-- Reads the configuration from <https://opentelemetry.io/docs/specs/otel/configuration/sdk-environment-variables/#general-sdk-configuration the standard environment variables>.
-- Delegates to 'runNoTracing' if @OTEL_SDK_DISABLED = true@.
runTracing
    :: ( IOE :> es
       , Concurrent :> es
       , Environment :> es
       , Retry :> es
       , Timeout :> es
       )
    => Resource
    -> Scope
    -> Eff (Tracing ': es) a
    -> Eff es a
runTracing resource scope eff =
    ifM Environment.isSdkDisabled (runNoTracing eff) do
        exporter <- Environment.runConfigError $ Exporter.Environment.lookup @Span resource scope
        runTracingWith exporter eff

-- | Run the 'Tracing' effect, sending telemetry to a collector at the given 'URI'
-- over HTTP with the given 'Encoding'.
runHttpTracing
    :: (IOE :> es, Concurrent :> es, Retry :> es, Timeout :> es)
    => Resource
    -> Scope
    -> Encoding
    -> Compression
    -> URI
    -> Eff (Tracing ': es) a
    -> Eff es a
runHttpTracing resource scope encoding compression endpoint =
    runTracingWith $
        Exporter.http
            @Span
            resource
            scope
            encoding
            endpoint
            (Export.defaultConfig @Span)
            compression
            responseTimeoutDefault

-- | Run the 'Tracing' effect, sending telemetry to a gRPC collector.
runGrpcTracing
    :: (IOE :> es, Concurrent :> es, Retry :> es, Timeout :> es)
    => Resource
    -> Scope
    -> HostName
    -> PortNumber
    -> Compression
    -> Eff (Tracing ': es) a
    -> Eff es a
runGrpcTracing resource scope host port compression =
    runTracingWith $
        Exporter.grpc @Span
            resource
            scope
            host
            port
            (Export.exportGrpcRPC @Span)
            (Export.defaultConfig @Span)
            compression
            defaultGrpcSendTimeout

-- | Run the 'Tracing' effect, printing telemetry to the console rather than sending to a collector.
runConsoleTracing :: (IOE :> es) => Eff (Tracing ': es) a -> Eff es a
runConsoleTracing = runTracingWith Console.stdout

-- | Run the 'Tracing' effect, collecting telemetry in-memory rather than sending to a collector.
runInMemoryTracing
    :: (IOE :> es, Concurrent :> es)
    => Eff (Tracing ': es) a
    -> Eff es (a, [Span])
runInMemoryTracing = runInMemoryOTLP @Span . runTracingState . inject

-- | Run the 'Tracing' effect with a custom 'Exporter'.
-- Passes 'Span's to the sink synchronously as they are finalised, without batching or retrying.
runTracingWith
    :: (IOE :> es)
    => Exporter es Span
    -> Eff (Tracing ': es) a
    -> Eff es a
runTracingWith exporter = runOTLPWith exporter . runTracingState . inject

-- | Run the 'Tracing' effect as a no-op action.
runNoTracing
    :: (IOE :> es)
    => Eff (Tracing ': es) a
    -> Eff es a
runNoTracing = runNoOTLP @Span . runTracingState . inject