-- | This module provides convenience functionality to debug traces locally. For production use,
-- prefer alternatives, e.g. "Monitor.Tracing.Zipkin".
module Monitor.Tracing.Local (
  collectSpanSamples
) where

import Control.Concurrent.STM (atomically, readTVar, readTChan, tryReadTChan)
import Control.Monad.Fix (fix)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trace
import Data.IORef (modifyIORef', newIORef, readIORef)
import UnliftIO (MonadUnliftIO)

-- | Runs a 'TraceT' action, returning any collected samples alongside its output. The samples are
-- sorted chronologically by completion time (e.g. the head is the first span to complete).
--
-- Spans which start before the action returns are guaranteed to be collected, even if they complete
-- after (in this case collection will block until their completion). More precisely,
-- 'collectSpanSamples' will return the first time there are no pending spans after the action is
-- done. For example:
--
-- > collectSpanSamples $ rootSpan alwaysSampled "parent" $ do
-- >   forkIO $ childSpan "child" $ threadDelay 2000000 -- Asynchronous 2 second child span.
-- >   threadDelay 1000000 -- Returns after one second, but the child span will still be sampled.
collectSpanSamples :: MonadUnliftIO m => TraceT m a -> m (a, [Sample])
collectSpanSamples :: TraceT m a -> m (a, [Sample])
collectSpanSamples TraceT m a
actn = do
  Tracer
tracer <- m Tracer
forall (m :: * -> *). MonadIO m => m Tracer
newTracer
  a
rv <- TraceT m a -> Tracer -> m a
forall (m :: * -> *) a. TraceT m a -> Tracer -> m a
runTraceT TraceT m a
actn Tracer
tracer
  IORef [Sample]
ref <- IO (IORef [Sample]) -> m (IORef [Sample])
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef [Sample]) -> m (IORef [Sample]))
-> IO (IORef [Sample]) -> m (IORef [Sample])
forall a b. (a -> b) -> a -> b
$ [Sample] -> IO (IORef [Sample])
forall a. a -> IO (IORef a)
newIORef []
  let
    addSample :: Sample -> m ()
addSample Sample
spl = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef [Sample] -> ([Sample] -> [Sample]) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef [Sample]
ref (Sample
splSample -> [Sample] -> [Sample]
forall a. a -> [a] -> [a]
:)
    samplesTC :: TChan Sample
samplesTC = Tracer -> TChan Sample
spanSamples Tracer
tracer
    pendingTV :: TVar Int
pendingTV = Tracer -> TVar Int
pendingSpanCount Tracer
tracer
  IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ (IO () -> IO ()) -> IO ()
forall a. (a -> a) -> a
fix ((IO () -> IO ()) -> IO ()) -> (IO () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \IO ()
loop -> do
    (Maybe Sample
mbSample, Int
pending) <- STM (Maybe Sample, Int) -> IO (Maybe Sample, Int)
forall a. STM a -> IO a
atomically (STM (Maybe Sample, Int) -> IO (Maybe Sample, Int))
-> STM (Maybe Sample, Int) -> IO (Maybe Sample, Int)
forall a b. (a -> b) -> a -> b
$ (,) (Maybe Sample -> Int -> (Maybe Sample, Int))
-> STM (Maybe Sample) -> STM (Int -> (Maybe Sample, Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TChan Sample -> STM (Maybe Sample)
forall a. TChan a -> STM (Maybe a)
tryReadTChan TChan Sample
samplesTC STM (Int -> (Maybe Sample, Int))
-> STM Int -> STM (Maybe Sample, Int)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
pendingTV
    case Maybe Sample
mbSample of
      Just Sample
spl -> Sample -> IO ()
forall (m :: * -> *). MonadIO m => Sample -> m ()
addSample Sample
spl IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
loop
      Maybe Sample
Nothing | Int
pending Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> IO Sample -> IO Sample
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (STM Sample -> IO Sample
forall a. STM a -> IO a
atomically (STM Sample -> IO Sample) -> STM Sample -> IO Sample
forall a b. (a -> b) -> a -> b
$ TChan Sample -> STM Sample
forall a. TChan a -> STM a
readTChan TChan Sample
samplesTC) IO Sample -> (Sample -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Sample -> IO ()
forall (m :: * -> *). MonadIO m => Sample -> m ()
addSample IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
loop
      Maybe Sample
_ -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  [Sample]
spls <- [Sample] -> [Sample]
forall a. [a] -> [a]
reverse ([Sample] -> [Sample]) -> m [Sample] -> m [Sample]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO [Sample] -> m [Sample]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IORef [Sample] -> IO [Sample]
forall a. IORef a -> IO a
readIORef IORef [Sample]
ref)
  (a, [Sample]) -> m (a, [Sample])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
rv, [Sample]
spls)