{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

-- | This module exposes the generic 'MonadTrace' class.
module Control.Monad.Trace.Class (
  -- * Types
  Span(..), spanIsSampled, spanIsDebug,
  Context(..),
  TraceID(..), decodeTraceID, encodeTraceID,
  SpanID(..), decodeSpanID, encodeSpanID,
  Reference(..),

  -- * Generating traces
  -- ** Individual spans
  MonadTrace(..),
  Builder(..), Name, builder,
  -- ** Structured traces
  rootSpan, rootSpanWith, childSpan, childSpanWith,
  -- ** Sampling
  SamplingDecision(..),
  SamplingPolicy, alwaysSampled, neverSampled, sampledWithProbability, sampledWhen, debugEnabled,

  -- * Annotating traces
  -- | Note that not all annotation types are supported by all backends. For example Zipkin only
  -- supports string tags (refer to "Monitor.Tracing.Zipkin" for the full list of supported span
  -- metadata).
  Key, Value, tagDoubleValue, tagInt64Value, tagTextValue, logValue, logValueAt
) where

import Control.Monad.Trace.Internal

import Control.Monad.Except (ExceptT(..))
import Control.Monad.Identity (Identity(..))
import Control.Monad.Reader (ReaderT(..))
import qualified Control.Monad.RWS.Lazy as RWS.Lazy
import qualified Control.Monad.RWS.Strict as RWS.Strict
import qualified Control.Monad.State.Lazy as State.Lazy
import qualified Control.Monad.State.Strict as State.Strict
import Control.Monad.Trans.Class (MonadTrans, lift)
import qualified Control.Monad.Writer.Lazy as Writer.Lazy
import qualified Control.Monad.Writer.Strict as Writer.Strict
import qualified Data.Aeson as JSON
import Data.ByteString (ByteString)
import Data.Int (Int64)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Data.String (IsString(..))
import Data.Text (Text)
import qualified Data.Text as T
import Data.Time.Clock.POSIX (POSIXTime)
import System.Random (randomRIO)

-- | A monad capable of generating and modifying trace spans.
--
-- This package currently provides two instances of this class:
--
-- * 'Control.Monad.Trace.TraceT', which emits spans for each trace in 'IO' and is meant to be used
-- in production.
-- * 'Identity', where tracing is a no-op and allows testing traced functions without any overhead
-- or complex setup.
class Monad m => MonadTrace m where

  -- | Trace an action, wrapping it inside a new span. If the action isn't currently being traced,
  -- 'trace' should be a no-op. Otherwise, the new span should share the active span's trace ID,
  -- sampling decision, and baggages unless overridden by the input 'Builder'.
  trace :: Builder -> m a -> m a

  -- | Extracts the currently active span, or 'Nothing' if the action is not being traced.
  activeSpan :: m (Maybe Span)
  default activeSpan :: (MonadTrace n, MonadTrans t, m ~ t n) => m (Maybe Span)
  activeSpan = n (Maybe Span) -> t n (Maybe Span)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift n (Maybe Span)
forall (m :: * -> *). MonadTrace m => m (Maybe Span)
activeSpan

  -- | Adds information to the active span, if present.
  addSpanEntry :: Key -> Value -> m ()
  default addSpanEntry :: (MonadTrace n, MonadTrans t, m ~ t n) => Key -> Value -> m ()
  addSpanEntry Key
key = n () -> t n ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (n () -> t n ()) -> (Value -> n ()) -> Value -> t n ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key -> Value -> n ()
forall (m :: * -> *). MonadTrace m => Key -> Value -> m ()
addSpanEntry Key
key

instance MonadTrace m => MonadTrace (ExceptT e m) where
  trace :: Builder -> ExceptT e m a -> ExceptT e m a
trace Builder
name (ExceptT m (Either e a)
actn) = m (Either e a) -> ExceptT e m a
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT (m (Either e a) -> ExceptT e m a)
-> m (Either e a) -> ExceptT e m a
forall a b. (a -> b) -> a -> b
$ Builder -> m (Either e a) -> m (Either e a)
forall (m :: * -> *) a. MonadTrace m => Builder -> m a -> m a
trace Builder
name m (Either e a)
actn

instance MonadTrace m => MonadTrace (ReaderT r m) where
  trace :: Builder -> ReaderT r m a -> ReaderT r m a
trace Builder
name (ReaderT r -> m a
actn) = (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m a) -> ReaderT r m a) -> (r -> m a) -> ReaderT r m a
forall a b. (a -> b) -> a -> b
$ \r
r -> Builder -> m a -> m a
forall (m :: * -> *) a. MonadTrace m => Builder -> m a -> m a
trace Builder
name (r -> m a
actn r
r)

instance (MonadTrace m, Monoid w) => MonadTrace (RWS.Lazy.RWST r w s m) where
  trace :: Builder -> RWST r w s m a -> RWST r w s m a
trace Builder
name (RWS.Lazy.RWST r -> s -> m (a, s, w)
actn) = (r -> s -> m (a, s, w)) -> RWST r w s m a
forall r w s (m :: * -> *) a.
(r -> s -> m (a, s, w)) -> RWST r w s m a
RWS.Lazy.RWST ((r -> s -> m (a, s, w)) -> RWST r w s m a)
-> (r -> s -> m (a, s, w)) -> RWST r w s m a
forall a b. (a -> b) -> a -> b
$ \r
r s
s -> Builder -> m (a, s, w) -> m (a, s, w)
forall (m :: * -> *) a. MonadTrace m => Builder -> m a -> m a
trace Builder
name (r -> s -> m (a, s, w)
actn r
r s
s)

instance (MonadTrace m, Monoid w) => MonadTrace (RWS.Strict.RWST r w s m) where
  trace :: Builder -> RWST r w s m a -> RWST r w s m a
trace Builder
name (RWS.Strict.RWST r -> s -> m (a, s, w)
actn) = (r -> s -> m (a, s, w)) -> RWST r w s m a
forall r w s (m :: * -> *) a.
(r -> s -> m (a, s, w)) -> RWST r w s m a
RWS.Strict.RWST ((r -> s -> m (a, s, w)) -> RWST r w s m a)
-> (r -> s -> m (a, s, w)) -> RWST r w s m a
forall a b. (a -> b) -> a -> b
$ \r
r s
s -> Builder -> m (a, s, w) -> m (a, s, w)
forall (m :: * -> *) a. MonadTrace m => Builder -> m a -> m a
trace Builder
name (r -> s -> m (a, s, w)
actn r
r s
s)

instance MonadTrace m => MonadTrace (State.Lazy.StateT s m) where
  trace :: Builder -> StateT s m a -> StateT s m a
trace Builder
name (State.Lazy.StateT s -> m (a, s)
actn) = (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
State.Lazy.StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$ \s
s -> Builder -> m (a, s) -> m (a, s)
forall (m :: * -> *) a. MonadTrace m => Builder -> m a -> m a
trace Builder
name (s -> m (a, s)
actn s
s)

instance MonadTrace m => MonadTrace (State.Strict.StateT s m) where
  trace :: Builder -> StateT s m a -> StateT s m a
trace Builder
name (State.Strict.StateT s -> m (a, s)
actn) = (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
State.Strict.StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$ \s
s -> Builder -> m (a, s) -> m (a, s)
forall (m :: * -> *) a. MonadTrace m => Builder -> m a -> m a
trace Builder
name (s -> m (a, s)
actn s
s)

instance (MonadTrace m, Monoid w) => MonadTrace (Writer.Lazy.WriterT w m) where
  trace :: Builder -> WriterT w m a -> WriterT w m a
trace Builder
name (Writer.Lazy.WriterT m (a, w)
actn) = m (a, w) -> WriterT w m a
forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
Writer.Lazy.WriterT (m (a, w) -> WriterT w m a) -> m (a, w) -> WriterT w m a
forall a b. (a -> b) -> a -> b
$ Builder -> m (a, w) -> m (a, w)
forall (m :: * -> *) a. MonadTrace m => Builder -> m a -> m a
trace Builder
name m (a, w)
actn

instance (MonadTrace m, Monoid w) => MonadTrace (Writer.Strict.WriterT w m) where
  trace :: Builder -> WriterT w m a -> WriterT w m a
trace Builder
name (Writer.Strict.WriterT m (a, w)
actn) = m (a, w) -> WriterT w m a
forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
Writer.Strict.WriterT (m (a, w) -> WriterT w m a) -> m (a, w) -> WriterT w m a
forall a b. (a -> b) -> a -> b
$ Builder -> m (a, w) -> m (a, w)
forall (m :: * -> *) a. MonadTrace m => Builder -> m a -> m a
trace Builder
name m (a, w)
actn

instance MonadTrace Identity where
  trace :: Builder -> Identity a -> Identity a
trace Builder
_ = Identity a -> Identity a
forall a. a -> a
id
  activeSpan :: Identity (Maybe Span)
activeSpan = Maybe Span -> Identity (Maybe Span)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Span
forall a. Maybe a
Nothing
  addSpanEntry :: Key -> Value -> Identity ()
addSpanEntry Key
_ Value
_ = () -> Identity ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- Creating traces

-- | A span builder.
--
-- 'Builder' has an 'IsString' instance, producing a span with the given string as name, no
-- additional references, tags, or baggages. This allows convenient creation of spans via the
-- @OverloadedStrings@ pragma.
data Builder = Builder
  { Builder -> Key
builderName :: !Name
  -- ^ Name of the generated span.
  , Builder -> Maybe TraceID
builderTraceID :: !(Maybe TraceID)
  -- ^ The trace ID of the generated span. If unset, the active span's trace ID will be used if
  -- present, otherwise a new ID will be generated.
  , Builder -> Maybe SpanID
builderSpanID :: !(Maybe SpanID)
  -- ^ The ID of the generated span, otherwise the ID will be auto-generated.
  , Builder -> Set Reference
builderReferences :: !(Set Reference)
  -- ^ Span references.
  , Builder -> Map Key Value
builderTags :: !(Map Key JSON.Value)
  -- ^ Initial set of tags.
  , Builder -> Map Key ByteString
builderBaggages :: !(Map Key ByteString)
  -- ^ Span context baggages.
  , Builder -> Maybe SamplingPolicy
builderSamplingPolicy :: !(Maybe SamplingPolicy)
  -- ^ How the span should be sampled. If unset, the active's span sampling will be used if present,
  -- otherwise the span will not be sampled.
  }

-- | Returns a 'Builder' with the given input as name and all other fields empty.
builder :: Name -> Builder
builder :: Key -> Builder
builder Key
name = Key
-> Maybe TraceID
-> Maybe SpanID
-> Set Reference
-> Map Key Value
-> Map Key ByteString
-> Maybe SamplingPolicy
-> Builder
Builder Key
name Maybe TraceID
forall a. Maybe a
Nothing Maybe SpanID
forall a. Maybe a
Nothing Set Reference
forall a. Set a
Set.empty Map Key Value
forall k a. Map k a
Map.empty Map Key ByteString
forall k a. Map k a
Map.empty Maybe SamplingPolicy
forall a. Maybe a
Nothing

instance IsString Builder where
  fromString :: String -> Builder
fromString = Key -> Builder
builder (Key -> Builder) -> (String -> Key) -> String -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Key
T.pack

-- | An action to determine how a span should be sampled.
type SamplingPolicy = IO SamplingDecision

-- | Returns a 'SamplingPolicy' which always samples.
alwaysSampled :: SamplingPolicy
alwaysSampled :: SamplingPolicy
alwaysSampled = SamplingDecision -> SamplingPolicy
forall (f :: * -> *) a. Applicative f => a -> f a
pure SamplingDecision
Always

-- | Returns a 'SamplingPolicy' which never samples.
neverSampled :: SamplingPolicy
neverSampled :: SamplingPolicy
neverSampled = SamplingDecision -> SamplingPolicy
forall (f :: * -> *) a. Applicative f => a -> f a
pure SamplingDecision
Never

-- | Returns a debug 'SamplingPolicy'. Debug spans are always sampled.
debugEnabled :: SamplingPolicy
debugEnabled :: SamplingPolicy
debugEnabled = SamplingDecision -> SamplingPolicy
forall (f :: * -> *) a. Applicative f => a -> f a
pure SamplingDecision
Debug

-- | Returns a 'SamplingPolicy' which samples a span iff the input is 'True'. It is equivalent to:
--
-- > sampledWhen b = if b then alwaysSampled else neverSampled
sampledWhen :: Bool -> SamplingPolicy
sampledWhen :: Bool -> SamplingPolicy
sampledWhen Bool
b = SamplingDecision -> SamplingPolicy
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SamplingDecision -> SamplingPolicy)
-> SamplingDecision -> SamplingPolicy
forall a b. (a -> b) -> a -> b
$ if Bool
b then SamplingDecision
Always else SamplingDecision
Never

-- | Returns a 'SamplingPolicy' which randomly samples spans.
sampledWithProbability :: Double -> SamplingPolicy
sampledWithProbability :: Double -> SamplingPolicy
sampledWithProbability Double
r = (Double, Double) -> IO Double
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Double
0, Double
1) IO Double -> (Double -> SamplingPolicy) -> SamplingPolicy
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> SamplingPolicy
sampledWhen (Bool -> SamplingPolicy)
-> (Double -> Bool) -> Double -> SamplingPolicy
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
r)

-- Generic span creation

-- | Starts a new trace, customizing the span builder. Note that the sampling input will override
-- any sampling customization set on the builder.
rootSpanWith :: MonadTrace m => (Builder -> Builder) -> SamplingPolicy -> Name -> m a -> m a
rootSpanWith :: (Builder -> Builder) -> SamplingPolicy -> Key -> m a -> m a
rootSpanWith Builder -> Builder
f SamplingPolicy
policy Key
name = Builder -> m a -> m a
forall (m :: * -> *) a. MonadTrace m => Builder -> m a -> m a
trace (Builder -> m a -> m a) -> Builder -> m a -> m a
forall a b. (a -> b) -> a -> b
$ (Builder -> Builder
f (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ Key -> Builder
builder Key
name) { builderSamplingPolicy :: Maybe SamplingPolicy
builderSamplingPolicy = SamplingPolicy -> Maybe SamplingPolicy
forall a. a -> Maybe a
Just SamplingPolicy
policy }

-- | Starts a new trace. For performance reasons, it is possible to customize how frequently tracing
-- information is collected. This allows fine-grain control on the overhead induced by tracing. For
-- example, you might only want to sample 1% of a very actively used call-path with
-- @sampledWithProbability 0.01@.
rootSpan :: MonadTrace m => SamplingPolicy -> Name -> m a -> m a
rootSpan :: SamplingPolicy -> Key -> m a -> m a
rootSpan = (Builder -> Builder) -> SamplingPolicy -> Key -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
(Builder -> Builder) -> SamplingPolicy -> Key -> m a -> m a
rootSpanWith Builder -> Builder
forall a. a -> a
id

-- | Extends a trace, same as 'childSpan' but also customizing the builder.
childSpanWith :: MonadTrace m => (Builder -> Builder) -> Name -> m a -> m a
childSpanWith :: (Builder -> Builder) -> Key -> m a -> m a
childSpanWith Builder -> Builder
f Key
name m a
actn = m (Maybe Span)
forall (m :: * -> *). MonadTrace m => m (Maybe Span)
activeSpan m (Maybe Span) -> (Maybe Span -> m a) -> m a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  Maybe Span
Nothing -> m a
actn
  Just Span
spn -> do
    let
      ctx :: Context
ctx = Span -> Context
spanContext Span
spn
      bldr :: Builder
bldr = Builder -> Builder
f (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ Key -> Builder
builder Key
name
      bldr' :: Builder
bldr' = Builder
bldr
        { builderTraceID :: Maybe TraceID
builderTraceID = TraceID -> Maybe TraceID
forall a. a -> Maybe a
Just (TraceID -> Maybe TraceID) -> TraceID -> Maybe TraceID
forall a b. (a -> b) -> a -> b
$ Context -> TraceID
contextTraceID Context
ctx
        , builderReferences :: Set Reference
builderReferences = Reference -> Set Reference -> Set Reference
forall a. Ord a => a -> Set a -> Set a
Set.insert (SpanID -> Reference
ChildOf (SpanID -> Reference) -> SpanID -> Reference
forall a b. (a -> b) -> a -> b
$ Context -> SpanID
contextSpanID Context
ctx) (Builder -> Set Reference
builderReferences Builder
bldr) }
    Builder -> m a -> m a
forall (m :: * -> *) a. MonadTrace m => Builder -> m a -> m a
trace Builder
bldr' m a
actn

-- | Extends a trace: the active span's ID will be added as a reference to a newly created span and
-- both spans will share the same trace ID. If no span is active, 'childSpan' is a no-op.
childSpan :: MonadTrace m => Name -> m a -> m a
childSpan :: Key -> m a -> m a
childSpan = (Builder -> Builder) -> Key -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
(Builder -> Builder) -> Key -> m a -> m a
childSpanWith Builder -> Builder
forall a. a -> a
id

-- Writing metadata

-- | Generates a tag value from a double.
tagDoubleValue :: Double -> Value
tagDoubleValue :: Double -> Value
tagDoubleValue = Value -> Value
TagValue (Value -> Value) -> (Double -> Value) -> Double -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Value
forall a. ToJSON a => a -> Value
JSON.toJSON

-- | Generates a 64-bit integer tag value from any integer.
tagInt64Value :: Integral a => a -> Value
tagInt64Value :: a -> Value
tagInt64Value = Value -> Value
TagValue (Value -> Value) -> (a -> Value) -> a -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ToJSON Int64 => Int64 -> Value
forall a. ToJSON a => a -> Value
JSON.toJSON @Int64) (Int64 -> Value) -> (a -> Int64) -> a -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- | Generates a Unicode text tag value.
tagTextValue :: Text -> Value
tagTextValue :: Key -> Value
tagTextValue = Value -> Value
TagValue (Value -> Value) -> (Key -> Value) -> Key -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key -> Value
forall a. ToJSON a => a -> Value
JSON.toJSON

-- | Generates a log value with the time of writing as timestamp. Note that the value may be written
-- later than it is created. For more control on the timestamp, use 'logValueAt'.
logValue :: JSON.ToJSON a => a -> Value
logValue :: a -> Value
logValue a
v = Value -> Maybe POSIXTime -> Value
LogValue (a -> Value
forall a. ToJSON a => a -> Value
JSON.toJSON a
v) Maybe POSIXTime
forall a. Maybe a
Nothing

-- | Generates a log value with a custom time.
logValueAt :: JSON.ToJSON a => POSIXTime -> a -> Value
logValueAt :: POSIXTime -> a -> Value
logValueAt POSIXTime
t a
v = Value -> Maybe POSIXTime -> Value
LogValue (a -> Value
forall a. ToJSON a => a -> Value
JSON.toJSON a
v) (POSIXTime -> Maybe POSIXTime
forall a. a -> Maybe a
Just POSIXTime
t)