{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
module Control.Monad.Trace (
Tracer, newTracer,
runTraceT, TraceT,
spanSamples, Sample(..), Tags, Logs,
pendingSpanCount,
) where
import Prelude hiding (span)
import Control.Monad.Trace.Class
import Control.Monad.Trace.Internal
import Control.Applicative ((<|>))
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Reader (ReaderT(..), ask, asks, local, runReaderT)
import Control.Monad.Reader.Class (MonadReader)
import Control.Monad.Trans.Class (MonadTrans(..))
import qualified Data.Aeson as JSON
import Data.List (sortOn)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe)
import Data.Time.Clock (NominalDiffTime)
import Data.Time.Clock.POSIX (POSIXTime, getPOSIXTime)
import UnliftIO (MonadUnliftIO, UnliftIO(..), askUnliftIO, withUnliftIO)
import UnliftIO.Exception (finally)
import UnliftIO.STM (TChan, TVar, atomically, modifyTVar', newTChanIO, newTVarIO, readTVar, writeTChan, writeTVar)
type Tags = Map Key JSON.Value
type Logs = [(POSIXTime, Key, JSON.Value)]
data Sample = Sample
{ sampleSpan :: !Span
, sampleTags :: !Tags
, sampleLogs :: !Logs
, sampleStart :: !POSIXTime
, sampleDuration :: !NominalDiffTime
}
data Tracer = Tracer
{ tracerChannel :: TChan Sample
, tracerPendingCount :: TVar Int
}
newTracer :: MonadIO m => m Tracer
newTracer = liftIO $ Tracer <$> newTChanIO <*> newTVarIO 0
pendingSpanCount :: Tracer -> TVar Int
pendingSpanCount = tracerPendingCount
spanSamples :: Tracer -> TChan Sample
spanSamples = tracerChannel
data Scope = Scope
{ scopeTracer :: !Tracer
, scopeSpan :: !(Maybe Span)
, scopeTags :: !(Maybe (TVar Tags))
, scopeLogs :: !(Maybe (TVar Logs))
}
newtype TraceT m a = TraceT { traceTReader :: ReaderT Scope m a }
deriving (Functor, Applicative, Monad, MonadIO, MonadTrans)
instance MonadReader r m => MonadReader r (TraceT m) where
ask = lift ask
local f (TraceT (ReaderT g)) = TraceT $ ReaderT $ \r -> local f $ g r
instance MonadUnliftIO m => MonadTrace (TraceT m) where
trace bldr (TraceT reader) = TraceT $ do
parentScope <- ask
let
mbParentSpn = scopeSpan parentScope
mbParentCtx = spanContext <$> mbParentSpn
mbTraceID = contextTraceID <$> mbParentCtx
spanID <- maybe (liftIO randomSpanID) pure $ builderSpanID bldr
traceID <- maybe (liftIO randomTraceID) pure $ builderTraceID bldr <|> mbTraceID
sampling <- case builderSamplingPolicy bldr of
Just policy -> liftIO policy
Nothing -> pure $ fromMaybe Never (spanSamplingDecision <$> mbParentSpn)
let
baggages = fromMaybe Map.empty $ contextBaggages <$> mbParentCtx
ctx = Context traceID spanID (builderBaggages bldr `Map.union` baggages)
spn = Span (builderName bldr) ctx (builderReferences bldr) sampling
tracer = scopeTracer parentScope
if spanIsSampled spn
then do
tagsTV <- newTVarIO $ builderTags bldr
logsTV <- newTVarIO []
startTV <- newTVarIO Nothing
let
run = do
start <- liftIO $ getPOSIXTime
atomically $ do
writeTVar startTV (Just start)
modifyTVar' (tracerPendingCount tracer) (+1)
local (const $ Scope tracer (Just spn) (Just tagsTV) (Just logsTV)) reader
cleanup = do
end <- liftIO $ getPOSIXTime
atomically $ readTVar startTV >>= \case
Nothing -> pure ()
Just start -> do
modifyTVar' (tracerPendingCount tracer) (\n -> n - 1)
tags <- readTVar tagsTV
logs <- sortOn (\(t, k, _) -> (t, k)) <$> readTVar logsTV
writeTChan (tracerChannel tracer) (Sample spn tags logs start (end - start))
run `finally` cleanup
else local (const $ Scope tracer (Just spn) Nothing Nothing) reader
activeSpan = TraceT $ asks scopeSpan
addSpanEntry key (TagValue val) = TraceT $ asks scopeTags >>= \case
Nothing -> pure ()
Just tv -> atomically $ modifyTVar' tv $ Map.insert key val
addSpanEntry key (LogValue val maybeTime) = TraceT $ asks scopeLogs >>= \case
Nothing -> pure ()
Just tv -> do
time <- case maybeTime of
Nothing -> liftIO getPOSIXTime
Just time' -> pure time'
atomically $ modifyTVar' tv ((time, key, val) :)
instance MonadUnliftIO m => MonadUnliftIO (TraceT m) where
askUnliftIO = TraceT $ withUnliftIO $ \u -> pure (UnliftIO (unliftIO u . traceTReader ))
runTraceT :: TraceT m a -> Tracer -> m a
runTraceT (TraceT reader) tracer = runReaderT reader (Scope tracer Nothing Nothing Nothing)