{-|
Module: OpenTracing.Standard

Standard implementations of `OpenTracing.Tracer` fields.
-}

{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns        #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE StrictData            #-}
{-# LANGUAGE TemplateHaskell       #-}

module OpenTracing.Standard
    ( StdEnv
    , newStdEnv
    , envTraceID128bit
    , envSampler

    , stdTracer
    , stdReporter
    )
where

import Control.Concurrent (MVar, newMVar, withMVar)
import Control.Lens                 hiding (Context, (.=))
import Control.Monad.Reader
import Data.Monoid
import Data.Word
import OpenTracing.Reporting.Stdio (stdoutReporter)
import OpenTracing.Sampling         (Sampler (runSampler))
import OpenTracing.Span
import OpenTracing.Types
import Prelude                      hiding (putStrLn)
import System.Random.MWC

-- | A standard environment for generating trace and span IDs.
data StdEnv = StdEnv
    { StdEnv -> MVar GenIO
envPRNGRef        :: MVar GenIO
    , StdEnv -> Sampler
_envSampler       :: Sampler
    , StdEnv -> Bool
_envTraceID128bit :: Bool
    }

newStdEnv :: MonadIO m => Sampler -> m StdEnv
newStdEnv :: Sampler -> m StdEnv
newStdEnv Sampler
samp = do
    Gen RealWorld
prng <- IO (Gen RealWorld) -> m (Gen RealWorld)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (Gen RealWorld)
IO GenIO
createSystemRandom
    MVar (Gen RealWorld)
prngRef <- IO (MVar (Gen RealWorld)) -> m (MVar (Gen RealWorld))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MVar (Gen RealWorld)) -> m (MVar (Gen RealWorld)))
-> IO (MVar (Gen RealWorld)) -> m (MVar (Gen RealWorld))
forall a b. (a -> b) -> a -> b
$ Gen RealWorld -> IO (MVar (Gen RealWorld))
forall a. a -> IO (MVar a)
newMVar Gen RealWorld
prng
    StdEnv -> m StdEnv
forall (m :: * -> *) a. Monad m => a -> m a
return StdEnv :: MVar GenIO -> Sampler -> Bool -> StdEnv
StdEnv { envPRNGRef :: MVar GenIO
envPRNGRef = MVar (Gen RealWorld)
MVar GenIO
prngRef, _envSampler :: Sampler
_envSampler = Sampler
samp, _envTraceID128bit :: Bool
_envTraceID128bit = Bool
True }

makeLenses ''StdEnv

-- | A standard implementation of `OpenTracing.Tracer.tracerStart`.
stdTracer :: MonadIO m => StdEnv -> SpanOpts -> m Span
stdTracer :: StdEnv -> SpanOpts -> m Span
stdTracer StdEnv
r = (ReaderT StdEnv m Span -> StdEnv -> m Span)
-> StdEnv -> ReaderT StdEnv m Span -> m Span
forall a b c. (a -> b -> c) -> b -> a -> c
flip ReaderT StdEnv m Span -> StdEnv -> m Span
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT StdEnv
r (ReaderT StdEnv m Span -> m Span)
-> (SpanOpts -> ReaderT StdEnv m Span) -> SpanOpts -> m Span
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SpanOpts -> ReaderT StdEnv m Span
forall (m :: * -> *).
(MonadIO m, MonadReader StdEnv m) =>
SpanOpts -> m Span
start

-- | A implementation of `OpenTracing.Tracer.tracerReport` that logs spans to stdout.
stdReporter :: MonadIO m => FinishedSpan -> m ()
stdReporter :: FinishedSpan -> m ()
stdReporter = FinishedSpan -> m ()
forall (m :: * -> *). MonadIO m => FinishedSpan -> m ()
stdoutReporter

--------------------------------------------------------------------------------
-- Internal

start :: (MonadIO m, MonadReader StdEnv m) => SpanOpts -> m Span
start :: SpanOpts -> m Span
start SpanOpts
so = do
    SpanContext
ctx <- do
        Maybe Reference
p <- [Reference] -> Maybe Reference
forall (t :: * -> *). Foldable t => t Reference -> Maybe Reference
findParent ([Reference] -> Maybe Reference)
-> m [Reference] -> m (Maybe Reference)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO [Reference] -> m [Reference]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (SpanRefs -> IO [Reference]
freezeRefs (Getting SpanRefs SpanOpts SpanRefs -> SpanOpts -> SpanRefs
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting SpanRefs SpanOpts SpanRefs
Lens' SpanOpts SpanRefs
spanOptRefs SpanOpts
so))
        case Maybe Reference
p of
            Maybe Reference
Nothing -> SpanOpts -> m SpanContext
forall (m :: * -> *).
(MonadIO m, MonadReader StdEnv m) =>
SpanOpts -> m SpanContext
freshContext SpanOpts
so
            Just Reference
p' -> SpanContext -> m SpanContext
forall (m :: * -> *).
(MonadIO m, MonadReader StdEnv m) =>
SpanContext -> m SpanContext
fromParent   (Reference -> SpanContext
refCtx Reference
p')
    SpanContext -> Text -> SpanRefs -> [Tag] -> m Span
forall (m :: * -> *) (t :: * -> *).
(MonadIO m, Foldable t) =>
SpanContext -> Text -> SpanRefs -> t Tag -> m Span
newSpan SpanContext
ctx
            (Getting Text SpanOpts Text -> SpanOpts -> Text
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Text SpanOpts Text
Lens' SpanOpts Text
spanOptOperation SpanOpts
so)
            (Getting SpanRefs SpanOpts SpanRefs -> SpanOpts -> SpanRefs
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting SpanRefs SpanOpts SpanRefs
Lens' SpanOpts SpanRefs
spanOptRefs SpanOpts
so)
            (Getting [Tag] SpanOpts [Tag] -> SpanOpts -> [Tag]
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting [Tag] SpanOpts [Tag]
Lens' SpanOpts [Tag]
spanOptTags SpanOpts
so)

newTraceID :: (MonadIO m, MonadReader StdEnv m) => m TraceID
newTraceID :: m TraceID
newTraceID = do
    StdEnv{Bool
MVar GenIO
Sampler
_envTraceID128bit :: Bool
_envSampler :: Sampler
envPRNGRef :: MVar GenIO
_envTraceID128bit :: StdEnv -> Bool
_envSampler :: StdEnv -> Sampler
envPRNGRef :: StdEnv -> MVar GenIO
..} <- m StdEnv
forall r (m :: * -> *). MonadReader r m => m r
ask
    IO TraceID -> m TraceID
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO TraceID -> m TraceID) -> IO TraceID -> m TraceID
forall a b. (a -> b) -> a -> b
$ MVar (Gen RealWorld) -> (Gen RealWorld -> IO TraceID) -> IO TraceID
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar (Gen RealWorld)
MVar GenIO
envPRNGRef ((Gen RealWorld -> IO TraceID) -> IO TraceID)
-> (Gen RealWorld -> IO TraceID) -> IO TraceID
forall a b. (a -> b) -> a -> b
$ \Gen RealWorld
prng -> do
      Maybe Word64
hi <- if Bool
_envTraceID128bit then
                Word64 -> Maybe Word64
forall a. a -> Maybe a
Just (Word64 -> Maybe Word64) -> IO Word64 -> IO (Maybe Word64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Word64 -> IO Word64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (GenIO -> IO Word64
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
Gen (PrimState m) -> m a
uniform Gen RealWorld
GenIO
prng)
            else
                Maybe Word64 -> IO (Maybe Word64)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Word64
forall a. Maybe a
Nothing
      Word64
lo <- IO Word64 -> IO Word64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Word64 -> IO Word64) -> IO Word64 -> IO Word64
forall a b. (a -> b) -> a -> b
$ GenIO -> IO Word64
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
Gen (PrimState m) -> m a
uniform Gen RealWorld
GenIO
prng
      TraceID -> IO TraceID
forall (m :: * -> *) a. Monad m => a -> m a
return TraceID :: Maybe Word64 -> Word64 -> TraceID
TraceID { traceIdHi :: Maybe Word64
traceIdHi = Maybe Word64
hi, traceIdLo :: Word64
traceIdLo = Word64
lo }

newSpanID :: (MonadIO m, MonadReader StdEnv m) => m Word64
newSpanID :: m Word64
newSpanID = do
  MVar (Gen RealWorld)
prngRef <- (StdEnv -> MVar (Gen RealWorld)) -> m (MVar (Gen RealWorld))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks StdEnv -> MVar (Gen RealWorld)
StdEnv -> MVar GenIO
envPRNGRef
  IO Word64 -> m Word64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Word64 -> m Word64) -> IO Word64 -> m Word64
forall a b. (a -> b) -> a -> b
$ MVar (Gen RealWorld) -> (Gen RealWorld -> IO Word64) -> IO Word64
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar (Gen RealWorld)
prngRef ((Gen RealWorld -> IO Word64) -> IO Word64)
-> (Gen RealWorld -> IO Word64) -> IO Word64
forall a b. (a -> b) -> a -> b
$ \Gen RealWorld
prng -> do
    GenIO -> IO Word64
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
Gen (PrimState m) -> m a
uniform Gen RealWorld
GenIO
prng

freshContext
    :: ( MonadIO            m
       , MonadReader StdEnv m
       )
    => SpanOpts
    -> m SpanContext
freshContext :: SpanOpts -> m SpanContext
freshContext SpanOpts
so = do
    TraceID
trid <- m TraceID
forall (m :: * -> *).
(MonadIO m, MonadReader StdEnv m) =>
m TraceID
newTraceID
    Word64
spid <- m Word64
forall (m :: * -> *). (MonadIO m, MonadReader StdEnv m) => m Word64
newSpanID
    Sampler
smpl <- Getting Sampler StdEnv Sampler -> m Sampler
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Sampler StdEnv Sampler
Lens' StdEnv Sampler
envSampler

    Sampled
sampled' <- case Getting (Maybe Sampled) SpanOpts (Maybe Sampled)
-> SpanOpts -> Maybe Sampled
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting (Maybe Sampled) SpanOpts (Maybe Sampled)
Lens' SpanOpts (Maybe Sampled)
spanOptSampled SpanOpts
so of
        Maybe Sampled
Nothing -> Getting Sampled Bool Sampled -> Bool -> Sampled
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Sampled Bool Sampled
Iso' Bool Sampled
_IsSampled
               (Bool -> Sampled) -> m Bool -> m Sampled
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sampler -> TraceID -> Text -> m Bool
Sampler
-> forall (m :: * -> *). MonadIO m => TraceID -> Text -> m Bool
runSampler Sampler
smpl TraceID
trid (Getting Text SpanOpts Text -> SpanOpts -> Text
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Text SpanOpts Text
Lens' SpanOpts Text
spanOptOperation SpanOpts
so)
        Just Sampled
s  -> Sampled -> m Sampled
forall (f :: * -> *) a. Applicative f => a -> f a
pure Sampled
s

    SpanContext -> m SpanContext
forall (m :: * -> *) a. Monad m => a -> m a
return SpanContext :: TraceID
-> Word64
-> Maybe Word64
-> Sampled
-> HashMap Text Text
-> SpanContext
SpanContext
        { ctxTraceID :: TraceID
ctxTraceID      = TraceID
trid
        , ctxSpanID :: Word64
ctxSpanID       = Word64
spid
        , ctxParentSpanID :: Maybe Word64
ctxParentSpanID = Maybe Word64
forall a. Maybe a
Nothing
        , _ctxSampled :: Sampled
_ctxSampled     = Sampled
sampled'
        , _ctxBaggage :: HashMap Text Text
_ctxBaggage     = HashMap Text Text
forall a. Monoid a => a
mempty
        }

fromParent
    :: ( MonadIO            m
       , MonadReader StdEnv m
       )
    => SpanContext
    -> m SpanContext
fromParent :: SpanContext -> m SpanContext
fromParent SpanContext
p = do
    Word64
spid <- m Word64
forall (m :: * -> *). (MonadIO m, MonadReader StdEnv m) => m Word64
newSpanID
    SpanContext -> m SpanContext
forall (m :: * -> *) a. Monad m => a -> m a
return SpanContext :: TraceID
-> Word64
-> Maybe Word64
-> Sampled
-> HashMap Text Text
-> SpanContext
SpanContext
        { ctxTraceID :: TraceID
ctxTraceID      = SpanContext -> TraceID
ctxTraceID SpanContext
p
        , ctxSpanID :: Word64
ctxSpanID       = Word64
spid
        , ctxParentSpanID :: Maybe Word64
ctxParentSpanID = Word64 -> Maybe Word64
forall a. a -> Maybe a
Just (SpanContext -> Word64
ctxSpanID SpanContext
p)
        , _ctxSampled :: Sampled
_ctxSampled     = Getting Sampled SpanContext Sampled -> SpanContext -> Sampled
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Sampled SpanContext Sampled
Lens' SpanContext Sampled
ctxSampled SpanContext
p
        , _ctxBaggage :: HashMap Text Text
_ctxBaggage     = Getting (HashMap Text Text) SpanContext (HashMap Text Text)
-> SpanContext -> HashMap Text Text
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting (HashMap Text Text) SpanContext (HashMap Text Text)
Lens' SpanContext (HashMap Text Text)
ctxBaggage SpanContext
p
        }