-- | Application tracing via <https://opentelemetry.io/>
--
-- @
-- data App = App
--   { -- ...
--   , appTracer :: Tracer
--   }
--
-- instance 'HasTracer' App where
--   tracerL = lens appTracer $ \x y -> x { appTracer = y }
--
-- loadApp f = do
--   -- ...
--   'withTracerProvider' $ \tracerProvider -> do
--     let appTracer = 'makeTracer' tracerProvider "my-app" 'tracerOptions'
--     f App {..}
-- @
--
-- You may need to do this even if you don't plan to manually trace things, in
-- order to satisfy the 'MonadTracer' constraint required by functions like
-- 'runDB'. If you don't need this feature, and don't plan on running an
-- otel-collector, set @OTEL_TRACES_EXPORTER=none@ in the environment, which
-- makes all tracing a no-op.
--
-- In the future, it should be possible to use @OTEL_SDK_DISABLED@ for the same
-- purpose. See <https://github.com/iand675/hs-opentelemetry/issues/60>.
module Freckle.App.OpenTelemetry
  ( HasTracer (..)
  , Tracer

    -- * Effects
  , MonadTracer (..)
  , inSpan
  , defaultSpanArguments
  , serverSpanArguments
  , clientSpanArguments
  , producerSpanArguments
  , consumerSpanArguments

    -- * Querying
  , withTraceIdContext
  , getCurrentTraceId
  , getCurrentTraceIdAsDatadog
  , getCurrentSpanContext

    -- * Setup
  , withTracerProvider

    -- ** 'Tracer'
  , makeTracer
  , tracerOptions

    -- ** Utilities
  , byteStringToAttribute
  , attributeValueLimit
  ) where

import Freckle.App.Prelude

import Blammo.Logging (withThreadContext, (.=))
import Control.Monad.Catch (MonadMask)
import Data.ByteString (ByteString)
import qualified Data.Text as T
import Data.Text.Encoding (decodeUtf8With)
import Data.Text.Encoding.Error (lenientDecode)
import Data.Word (Word64)
import OpenTelemetry.Context (lookupSpan)
import OpenTelemetry.Context.ThreadLocal (getContext)
import OpenTelemetry.Propagator.Datadog
  ( convertOpenTelemetryTraceIdToDatadogTraceId
  )
import OpenTelemetry.Trace hiding (inSpan)
import OpenTelemetry.Trace.Core (getSpanContext)
import qualified OpenTelemetry.Trace.Core as Trace (SpanContext (..))
import OpenTelemetry.Trace.Id (TraceId)
import OpenTelemetry.Trace.Monad
import UnliftIO.Exception (bracket)

-- | 'defaultSpanArguments' with 'kind' set to 'Server'
--
-- Indicates that the span covers server-side handling of a synchronous RPC or
-- other remote request. This span is the child of a remote @Client@ span that
-- was expected to wait for a response.
serverSpanArguments :: SpanArguments
serverSpanArguments :: SpanArguments
serverSpanArguments = SpanArguments
defaultSpanArguments {kind = Server}

-- | 'defaultSpanArguments' with 'kind' set to 'Kind'
--
-- Indicates that the span describes a synchronous request to some remote
-- service. This span is the parent of a remote @Server@ span and waits for its
-- response.
clientSpanArguments :: SpanArguments
clientSpanArguments :: SpanArguments
clientSpanArguments = SpanArguments
defaultSpanArguments {kind = Client}

-- | 'defaultSpanArguments' with 'kind' set to 'Producer'
--
-- Indicates that the span describes the parent of an asynchronous request. This
-- parent span is expected to end before the corresponding child @Producer@
-- span, possibly even before the child span starts. In messaging scenarios with
-- batching, tracing individual messages requires a new @Producer@ span per
-- message to be created.
producerSpanArguments :: SpanArguments
producerSpanArguments :: SpanArguments
producerSpanArguments = SpanArguments
defaultSpanArguments {kind = Producer}

-- | 'defaultSpanArguments' with 'kind' set to 'Consumer'
--
-- Indicates that the span describes the child of an asynchronous @Producer@
-- request.
consumerSpanArguments :: SpanArguments
consumerSpanArguments :: SpanArguments
consumerSpanArguments = SpanArguments
defaultSpanArguments {kind = Consumer}

withTracerProvider :: MonadUnliftIO m => (TracerProvider -> m a) -> m a
withTracerProvider :: forall (m :: * -> *) a.
MonadUnliftIO m =>
(TracerProvider -> m a) -> m a
withTracerProvider =
  m TracerProvider
-> (TracerProvider -> m ()) -> (TracerProvider -> m a) -> m a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket
    (IO TracerProvider -> m TracerProvider
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO TracerProvider
initializeGlobalTracerProvider)
    (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ())
-> (TracerProvider -> IO ()) -> TracerProvider -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TracerProvider -> IO ()
forall (m :: * -> *). MonadIO m => TracerProvider -> m ()
shutdownTracerProvider)

getCurrentTraceId :: MonadIO m => m (Maybe TraceId)
getCurrentTraceId :: forall (m :: * -> *). MonadIO m => m (Maybe TraceId)
getCurrentTraceId = (SpanContext -> TraceId) -> Maybe SpanContext -> Maybe TraceId
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SpanContext -> TraceId
Trace.traceId (Maybe SpanContext -> Maybe TraceId)
-> m (Maybe SpanContext) -> m (Maybe TraceId)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Maybe SpanContext)
forall (m :: * -> *). MonadIO m => m (Maybe SpanContext)
getCurrentSpanContext

getCurrentTraceIdAsDatadog :: MonadIO m => m (Maybe Word64)
getCurrentTraceIdAsDatadog :: forall (m :: * -> *). MonadIO m => m (Maybe Word64)
getCurrentTraceIdAsDatadog =
  (TraceId -> Word64) -> Maybe TraceId -> Maybe Word64
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TraceId -> Word64
convertOpenTelemetryTraceIdToDatadogTraceId (Maybe TraceId -> Maybe Word64)
-> m (Maybe TraceId) -> m (Maybe Word64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Maybe TraceId)
forall (m :: * -> *). MonadIO m => m (Maybe TraceId)
getCurrentTraceId

getCurrentSpanContext :: MonadIO m => m (Maybe SpanContext)
getCurrentSpanContext :: forall (m :: * -> *). MonadIO m => m (Maybe SpanContext)
getCurrentSpanContext = do
  Maybe Span
mSpan <- Context -> Maybe Span
lookupSpan (Context -> Maybe Span) -> m Context -> m (Maybe Span)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Context
forall (m :: * -> *). MonadIO m => m Context
getContext
  (Span -> m SpanContext) -> Maybe Span -> m (Maybe SpanContext)
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) -> Maybe a -> f (Maybe b)
traverse Span -> m SpanContext
forall (m :: * -> *). MonadIO m => Span -> m SpanContext
getSpanContext Maybe Span
mSpan

withTraceIdContext :: (MonadIO m, MonadMask m) => m a -> m a
withTraceIdContext :: forall (m :: * -> *) a. (MonadIO m, MonadMask m) => m a -> m a
withTraceIdContext m a
f = do
  Maybe Word64
mTraceId <- m (Maybe Word64)
forall (m :: * -> *). MonadIO m => m (Maybe Word64)
getCurrentTraceIdAsDatadog
  m a -> (Word64 -> m a) -> Maybe Word64 -> m a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m a
f (\Word64
traceId -> [Pair] -> m a -> m a
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
[Pair] -> m a -> m a
withThreadContext [Key
"trace_id" Key -> Word64 -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= Word64
traceId] m a
f) Maybe Word64
mTraceId

-- | Convert a 'ByteString' to an 'Attribute' safely
--
-- - Decodes it as UTF-8 leniently,
-- - Truncates to fit within 'attributeValueLimit'
byteStringToAttribute :: ByteString -> Attribute
byteStringToAttribute :: ByteString -> Attribute
byteStringToAttribute =
  Text -> Attribute
forall a. ToAttribute a => a -> Attribute
toAttribute
    (Text -> Attribute)
-> (ByteString -> Text) -> ByteString -> Attribute
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Text -> Text
truncateText Int
attributeValueLimit
    (Text -> Text) -> (ByteString -> Text) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OnDecodeError -> ByteString -> Text
decodeUtf8With OnDecodeError
lenientDecode

-- | Character limit for 'Attribute' values
--
-- OTel the spec doesn't specify a limit, but says that SDKs should decide
-- some limit. It's not clear what the Haskell SDK does, if anything. New
-- Relic applies a limit of 4095 characters on all metrics it handles,
-- including those coming from OTel. Seems reasonable enough.
--
-- <https://docs.newrelic.com/docs/more-integrations/open-source-telemetry-integrations/opentelemetry/best-practices/opentelemetry-best-practices-attributes/>
attributeValueLimit :: Int
attributeValueLimit :: Int
attributeValueLimit = Int
4095

truncateText :: Int -> Text -> Text
truncateText :: Int -> Text -> Text
truncateText Int
l Text
t
  | Text -> Int
T.length Text
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l = Text
t
  | Bool
otherwise = Int -> Text -> Text
T.take (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
3) Text
t Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"..."