Allow discarding spans

src/Control/Monad/Trace.hs

@@ -9,7 +9,7 @@
 module Control.Monad.Trace (
   -- * Tracers
   Tracer, newTracer,
-  runTraceT, TraceT,
+  runTraceT, runTraceT', TraceT,
 
   -- * Collected data
   -- | Tracers currently expose two pieces of data: completed spans and pending span count. Note
@@ -101,61 +101,63 @@ data Scope = Scope
   }
 
 -- | A span generation monad.
-newtype TraceT m a = TraceT { traceTReader :: ReaderT Scope m a }
+newtype TraceT m a = TraceT { traceTReader :: ReaderT (Maybe 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 -- To detect whether an exception happened during span setup.
-        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 () -- The action was interrupted before the span was pending.
-              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
+  trace bldr (TraceT reader) = TraceT $ ask >>= \case
+    Nothing -> reader
+    Just parentScope -> do
+      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 -- To detect whether an exception happened during span setup.
+          let
+            scope = Scope tracer (Just spn) (Just tagsTV) (Just logsTV)
+            run = do
+              start <- liftIO $ getPOSIXTime
+              atomically $ do
+                writeTVar startTV (Just start)
+                modifyTVar' (tracerPendingCount tracer) (+1)
+              local (const $ Just scope) reader
+            cleanup = do
+              end <- liftIO $ getPOSIXTime
+              atomically $ readTVar startTV >>= \case
+                Nothing -> pure () -- The action was interrupted before the span was pending.
+                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 $ Just $ Scope tracer (Just spn) Nothing Nothing) reader
+
+  activeSpan = TraceT $ asks (>>= scopeSpan)
 
   addSpanEntry key (TagValue val) = TraceT $ do
-    mbTV <- asks scopeTags
+    mbTV <- asks (>>= scopeTags)
     for_ mbTV $ \tv -> atomically $ modifyTVar' tv $ Map.insert key val
   addSpanEntry key (LogValue val mbTime)  = TraceT $ do
-    mbTV <- asks scopeLogs
+    mbTV <- asks (>>= scopeLogs)
     for_ mbTV $ \tv -> do
       time <- maybe (liftIO getPOSIXTime) pure mbTime
       atomically $ modifyTVar' tv ((time, key, val) :)
@@ -170,5 +172,13 @@ instance MonadUnliftIO m => MonadUnliftIO (TraceT m) where
 -- method explicitly. Instead, prefer to use the backend's functionality directly (e.g.
 -- 'Monitor.Tracing.Zipkin.run' for Zipkin). To ease debugging in certain cases,
 -- 'Monitor.Tracing.Local.collectSpanSamples' is also available.
+--
+-- See 'runTraceT'' for a variant which allows discarding spans.
 runTraceT :: TraceT m a -> Tracer -> m a
-runTraceT (TraceT reader) tracer = runReaderT reader (Scope tracer Nothing Nothing Nothing)
+runTraceT actn tracer = runTraceT' actn (Just tracer)
+
+-- | Maybe trace an action. If the tracer is 'Nothing', no spans will be published.
+runTraceT' :: TraceT m a -> Maybe Tracer -> m a
+runTraceT' (TraceT reader) mbTracer =
+  let scope = fmap (\tracer -> Scope tracer Nothing Nothing Nothing) mbTracer
+  in runReaderT reader scope

test/Spec.hs

@@ -9,6 +9,7 @@ import Monitor.Tracing
 import Monitor.Tracing.Local (collectSpanSamples)
 import qualified Monitor.Tracing.Zipkin as ZPK
 
+import Control.Monad.IO.Class (liftIO, MonadIO)
 import Control.Monad (void)
 import Control.Monad.Reader (MonadReader, Reader, ReaderT, ask, runReader, runReaderT)
 import Control.Monad.State.Strict (MonadState, StateT, evalStateT, get)
@@ -36,26 +37,40 @@ main = hspec $ do
           pure $ s + r
         v = runReader (evalStateT actn 1) 2
       v `shouldBe` 3
+
+    it "should be runnable in IO without a tracer" $ do
+      let
+        actn :: (MonadIO m, MonadTrace m) => m Int
+        actn = trace "one" $ do
+          r <- liftIO $ newIORef 1
+          trace "two" $ liftIO (readIORef r)
+      v <- runTraceT' actn Nothing
+      v `shouldBe` 1
+
   describe "trace" $ do
     it "should not create spans when no traces are started" $ do
       spans <- collectSpans $ pure ()
       fmap spanName spans `shouldBe` []
+
     it "should collect a single span when no children are created" $ do
       spans <- collectSpans (trace "t" { builderSamplingPolicy = Just alwaysSampled } $ pure ())
       fmap spanName spans `shouldBe` ["t"]
+
     it "should be able to stack on top of a ReaderT" $ do
       let
         actn = trace "t" { builderSamplingPolicy = Just alwaysSampled } $ do
           name <- ask
           trace (builder name) $ pure ()
       spans <- runReaderT (collectSpans @(ReaderT Text IO) actn) "foo"
       fmap spanName spans `shouldBe` ["foo", "t"]
+
   describe "Zipkin" $ do
     it "should round-trip a B3 using a single header" $ do
       let
         bs = "80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-1-05e3ac9a4f6e3b90"
         mbBs = ZPK.b3ToHeaderValue <$> ZPK.b3FromHeaderValue bs
       mbBs `shouldBe` Just bs
+
     it "should have equivalent B3 header representations" $ do
       let
         bs = "80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-1-05e3ac9a4f6e3b90"
@@ -67,6 +82,7 @@ main = hspec $ do
         Just b3 = ZPK.b3FromHeaderValue bs
         Just b3' = ZPK.b3FromHeaders hdrs
       b3 `shouldBe` b3'
+
     it "consumerSpan should use B3 as parent reference" $ do
       let
         bs = "80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-1-05e3ac9a4f6e3b90"
@@ -75,6 +91,7 @@ main = hspec $ do
       contextTraceID (spanContext consumerSpan) `shouldBe` ZPK.b3TraceID b3            -- same traceId
       contextSpanID (spanContext consumerSpan) `shouldNotBe` ZPK.b3SpanID b3           -- different spanId
       spanReferences consumerSpan `shouldBe` Set.singleton (ChildOf $ ZPK.b3SpanID b3) -- b3 spanId is parent
+
   describe "collectSpanSamples" $ do
     it "should collect spans which are still pending after the action returns" $ do
       spans <- collectSpans $ rootSpan alwaysSampled "sleep-parent" $ do