lightstep-haskell-0.1.0: src/LightStep/WithSpan.hs

module LightStep.WithSpan where

import Data.Maybe
import Data.ProtoLens.Message (defMessage)
import Control.Lens
import Control.Concurrent
import Control.Monad.Catch
import qualified Data.Text as T
import LightStep.GlobalSharedMutableSingleton
import LightStep.LowLevel
import System.IO.Unsafe
import Proto.Collector
import Proto.Collector_Fields
import qualified Data.HashMap.Strict as HM

{-# noinline globalSharedMutableSpanStacks #-}
globalSharedMutableSpanStacks :: MVar (HM.HashMap ThreadId [Span])
globalSharedMutableSpanStacks = unsafePerformIO (newMVar mempty)

withSpan :: T.Text -> IO a -> IO a
withSpan opName action =
  bracket
    (pushSpan opName)
    popSpan 
    (const action)

pushSpan :: T.Text -> IO ()
pushSpan opName = do
  sp <- startSpan opName
  tId <- myThreadId
  modifyMVar_ globalSharedMutableSpanStacks $ \stacks ->
    case fromMaybe [] (HM.lookup tId stacks) of
      [] -> do
        pure $! HM.insert tId [sp] stacks
      (psp : _) ->
        let !sp' = sp
              & references .~ [defMessage & relationship .~ Reference'CHILD_OF & spanContext .~ (psp ^. spanContext)]
              & spanContext.traceId .~ (psp ^. spanContext.traceId)
        in pure $! HM.update (Just . (sp' :)) tId stacks

popSpan :: () -> IO ()
popSpan () = do
  tId <- myThreadId
  sp <- modifyMVar globalSharedMutableSpanStacks
    (\stacks ->
      let (sp : sps) = stacks HM.! tId
          !stacks' = HM.insert tId sps stacks
      in pure (stacks', sp))
  sp' <- finishSpan sp
  submitSpan sp'

modifyCurrentSpan :: (Span -> Span) -> IO ()
modifyCurrentSpan f = do
  tId <- myThreadId
  modifyMVar_ globalSharedMutableSpanStacks
    (\stacks ->
      let (sp : sps) = stacks HM.! tId
          !stacks' = HM.insert tId (f sp : sps) stacks
      in pure stacks')

setTag :: T.Text -> T.Text -> IO ()
setTag k v =
  modifyCurrentSpan (tags %~ (<> [defMessage & key .~ k & stringValue .~ v]))