{-# 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(ReaderT), ask, asks, local, runReaderT)
import Control.Monad.Reader.Class (MonadReader)
import Control.Monad.Trans.Class (MonadTrans, lift)
import qualified Data.Aeson as JSON
import Data.Foldable (for_)
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, withRunInIO)
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 $ do
mbTV <- asks scopeTags
for_ mbTV $ \tv -> atomically $ modifyTVar' tv $ Map.insert key val
addSpanEntry key (LogValue val mbTime) = TraceT $ do
mbTV <- asks scopeLogs
for_ mbTV $ \tv -> do
time <- maybe (liftIO getPOSIXTime) pure mbTime
atomically $ modifyTVar' tv ((time, key, val) :)
instance MonadUnliftIO m => MonadUnliftIO (TraceT m) where
withRunInIO inner = TraceT $ withRunInIO $ \run -> inner (run . traceTReader)
runTraceT :: TraceT m a -> Tracer -> m a
runTraceT (TraceT reader) tracer = runReaderT reader (Scope tracer Nothing Nothing Nothing)