{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}

-- | This module implements a <https://zipkin.apache.org/ Zipkin>-powered trace publisher. You will
-- almost certainly want to import it qualified.
--
-- Zipkin does not support all OpenTracing functionality. To guarantee that everything works as
-- expected, you should only use the functions defined in this module or exported by
-- "Monitor.Tracing".
module Monitor.Tracing.Zipkin (
  -- * Configuration
  -- ** General settings
  Settings(..), defaultSettings,
  -- ** Endpoint
  Endpoint(..), defaultEndpoint,

  -- * Publishing traces
  Zipkin,
  new, run, publish, with,

  -- * Cross-process spans
  -- ** Communication
  B3(..), b3ToHeaders, b3FromHeaders, b3ToHeaderValue, b3FromHeaderValue, b3FromSpan,
  -- ** Span generation
  clientSpan, clientSpanWith, serverSpan, serverSpanWith, producerSpanWith, consumerSpanWith,

  -- * Custom metadata
  -- ** Tags
  tag, addTag, addInheritedTag, addProducerKind,
  -- ** Annotations
  -- | Annotations are similar to tags, but timestamped.
  annotate, annotateAt,
  -- ** Endpoints
  addEndpoint
) where

import Control.Monad.Trace
import Control.Monad.Trace.Class

import Control.Concurrent (forkIO, threadDelay)
import Control.Concurrent.STM (atomically, tryReadTChan)
import Control.Monad (forever, guard, void, when)
import Control.Monad.Fix (fix)
import Control.Monad.IO.Class (MonadIO, liftIO)
import qualified Data.Aeson as JSON
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BS
import Data.CaseInsensitive (CI)
import Data.Time.Clock (NominalDiffTime)
import Data.Foldable (toList)
import Data.Int (Int64)
import Data.IORef (modifyIORef, newIORef, readIORef)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (catMaybes, fromMaybe, listToMaybe, maybeToList)
import Data.Monoid (Endo(..))
#if !MIN_VERSION_base(4, 11, 0)
import Data.Semigroup ((<>))
#endif
import Data.Set (Set)
import qualified Data.Set as Set
import Data.String (IsString(..))
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import Data.Time.Clock.POSIX (POSIXTime)
import Network.HTTP.Client (Manager, Request)
import qualified Network.HTTP.Client as HTTP
import Network.Socket (HostName, PortNumber)
import UnliftIO (MonadUnliftIO)
import UnliftIO.Exception (finally)

-- | 'Zipkin' creation settings.
data Settings = Settings
  { Settings -> Maybe HostName
settingsHostname :: !(Maybe HostName)
  -- ^ The Zipkin server's hostname, defaults to @localhost@ if unset.
  , Settings -> Maybe PortNumber
settingsPort :: !(Maybe PortNumber)
  -- ^ The port the Zipkin server is listening on, defaults to @9411@ if unset.
  , Settings -> Maybe Endpoint
settingsEndpoint :: !(Maybe Endpoint)
  -- ^ Local endpoint included in all published spans.
  , Settings -> Maybe Manager
settingsManager :: !(Maybe Manager)
  -- ^ An optional HTTP manager to use for publishing spans on the Zipkin server.
  , Settings -> Maybe NominalDiffTime
settingsPublishPeriod :: !(Maybe NominalDiffTime)
  -- ^ If set to a positive value, traces will be flushed in the background every such period.
  }

-- | Creates empty 'Settings'. You will typically use this (or the 'IsString' instance) as starting
-- point to only fill in the fields you care about:
--
-- > let settings = defaultSettings { settingsPort = Just 2222 }
defaultSettings :: Settings
defaultSettings :: Settings
defaultSettings = Maybe HostName
-> Maybe PortNumber
-> Maybe Endpoint
-> Maybe Manager
-> Maybe NominalDiffTime
-> Settings
Settings Maybe HostName
forall a. Maybe a
Nothing Maybe PortNumber
forall a. Maybe a
Nothing Maybe Endpoint
forall a. Maybe a
Nothing Maybe Manager
forall a. Maybe a
Nothing Maybe NominalDiffTime
forall a. Maybe a
Nothing

-- | Generates settings with the given string as hostname.
instance IsString Settings where
  fromString :: HostName -> Settings
fromString HostName
s = Settings
defaultSettings { settingsHostname :: Maybe HostName
settingsHostname = HostName -> Maybe HostName
forall a. a -> Maybe a
Just HostName
s }

-- | A Zipkin trace publisher.
--
-- All publisher functionality is thread-safe. In particular it is safe to 'publish' concurrently
-- with 'run', and/or 'run' multiple actions concurrently. Note also that all sampled spans are
-- retained in memory until they are published.
data Zipkin = Zipkin
  { Zipkin -> Manager
zipkinManager :: !Manager
  , Zipkin -> Request
zipkinRequest :: !Request
  , Zipkin -> Tracer
zipkinTracer :: !Tracer
  , Zipkin -> Maybe Endpoint
zipkinEndpoint :: !(Maybe Endpoint)
  }

flushSpans :: Maybe Endpoint -> Tracer -> Request -> Manager -> IO ()
flushSpans :: Maybe Endpoint -> Tracer -> Request -> Manager -> IO ()
flushSpans Maybe Endpoint
ept Tracer
tracer Request
req Manager
mgr = do
  IORef [ZipkinSpan]
ref <- [ZipkinSpan] -> IO (IORef [ZipkinSpan])
forall a. a -> IO (IORef a)
newIORef []
  (IO () -> IO ()) -> IO ()
forall a. (a -> a) -> a
fix ((IO () -> IO ()) -> IO ()) -> (IO () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \IO ()
loop -> STM (Maybe Sample) -> IO (Maybe Sample)
forall a. STM a -> IO a
atomically (TChan Sample -> STM (Maybe Sample)
forall a. TChan a -> STM (Maybe a)
tryReadTChan (TChan Sample -> STM (Maybe Sample))
-> TChan Sample -> STM (Maybe Sample)
forall a b. (a -> b) -> a -> b
$ Tracer -> TChan Sample
spanSamples Tracer
tracer) IO (Maybe Sample) -> (Maybe Sample -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe Sample
Nothing -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Just Sample
sample -> IORef [ZipkinSpan] -> ([ZipkinSpan] -> [ZipkinSpan]) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef IORef [ZipkinSpan]
ref (Maybe Endpoint -> Sample -> ZipkinSpan
ZipkinSpan Maybe Endpoint
ept Sample
sampleZipkinSpan -> [ZipkinSpan] -> [ZipkinSpan]
forall a. a -> [a] -> [a]
:) IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
loop
  [ZipkinSpan]
spns <- IORef [ZipkinSpan] -> IO [ZipkinSpan]
forall a. IORef a -> IO a
readIORef IORef [ZipkinSpan]
ref
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [ZipkinSpan] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [ZipkinSpan]
spns) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    let req' :: Request
req' = Request
req { requestBody :: RequestBody
HTTP.requestBody = ByteString -> RequestBody
HTTP.RequestBodyLBS (ByteString -> RequestBody) -> ByteString -> RequestBody
forall a b. (a -> b) -> a -> b
$ [ZipkinSpan] -> ByteString
forall a. ToJSON a => a -> ByteString
JSON.encode [ZipkinSpan]
spns }
    IO (Response ByteString) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Response ByteString) -> IO ())
-> IO (Response ByteString) -> IO ()
forall a b. (a -> b) -> a -> b
$ Request -> Manager -> IO (Response ByteString)
HTTP.httpLbs Request
req' Manager
mgr

-- | Creates a 'Zipkin' publisher for the input 'Settings'.
new :: MonadIO m => Settings -> m Zipkin
new :: Settings -> m Zipkin
new (Settings Maybe HostName
mbHostname Maybe PortNumber
mbPort Maybe Endpoint
mbEpt Maybe Manager
mbMgr Maybe NominalDiffTime
mbPrd) = IO Zipkin -> m Zipkin
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Zipkin -> m Zipkin) -> IO Zipkin -> m Zipkin
forall a b. (a -> b) -> a -> b
$ do
  Manager
mgr <- IO Manager
-> (Manager -> IO Manager) -> Maybe Manager -> IO Manager
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (ManagerSettings -> IO Manager
HTTP.newManager ManagerSettings
HTTP.defaultManagerSettings) Manager -> IO Manager
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Manager
mbMgr
  Tracer
tracer <- IO Tracer
forall (m :: * -> *). MonadIO m => m Tracer
newTracer
  let
    req :: Request
req = Request
HTTP.defaultRequest
      { method :: Method
HTTP.method = Method
"POST"
      , host :: Method
HTTP.host = HostName -> Method
BS.pack (HostName -> Maybe HostName -> HostName
forall a. a -> Maybe a -> a
fromMaybe HostName
"localhost" Maybe HostName
mbHostname)
      , path :: Method
HTTP.path = Method
"/api/v2/spans"
      , port :: Int
HTTP.port = Int -> (PortNumber -> Int) -> Maybe PortNumber -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
9411 PortNumber -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Maybe PortNumber
mbPort
      , requestHeaders :: RequestHeaders
HTTP.requestHeaders = [(HeaderName
"content-type", Method
"application/json")]
      }
  IO (Maybe ThreadId) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe ThreadId) -> IO ()) -> IO (Maybe ThreadId) -> IO ()
forall a b. (a -> b) -> a -> b
$ let prd :: NominalDiffTime
prd = NominalDiffTime -> Maybe NominalDiffTime -> NominalDiffTime
forall a. a -> Maybe a -> a
fromMaybe NominalDiffTime
0 Maybe NominalDiffTime
mbPrd in if NominalDiffTime
prd NominalDiffTime -> NominalDiffTime -> Bool
forall a. Ord a => a -> a -> Bool
<= NominalDiffTime
0
    then Maybe ThreadId -> IO (Maybe ThreadId)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ThreadId
forall a. Maybe a
Nothing
    else (ThreadId -> Maybe ThreadId) -> IO ThreadId -> IO (Maybe ThreadId)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ThreadId -> Maybe ThreadId
forall a. a -> Maybe a
Just (IO ThreadId -> IO (Maybe ThreadId))
-> IO ThreadId -> IO (Maybe ThreadId)
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      Int -> IO ()
threadDelay (NominalDiffTime -> Int
forall a. Integral a => NominalDiffTime -> a
microSeconds NominalDiffTime
prd)
      Maybe Endpoint -> Tracer -> Request -> Manager -> IO ()
flushSpans Maybe Endpoint
mbEpt Tracer
tracer Request
req Manager
mgr -- Manager is thread-safe.
  Zipkin -> IO Zipkin
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Zipkin -> IO Zipkin) -> Zipkin -> IO Zipkin
forall a b. (a -> b) -> a -> b
$ Manager -> Request -> Tracer -> Maybe Endpoint -> Zipkin
Zipkin Manager
mgr Request
req Tracer
tracer Maybe Endpoint
mbEpt

-- | Runs a 'TraceT' action, sampling spans appropriately. Note that this method does not publish
-- spans on its own; to do so, either call 'publish' manually or specify a positive
-- 'settingsPublishPeriod' to publish in the background.
run :: TraceT m a -> Zipkin -> m a
run :: TraceT m a -> Zipkin -> m a
run TraceT m a
actn Zipkin
zipkin = TraceT m a -> Tracer -> m a
forall (m :: * -> *) a. TraceT m a -> Tracer -> m a
runTraceT TraceT m a
actn (Zipkin -> Tracer
zipkinTracer Zipkin
zipkin)

-- | Flushes all complete spans to the Zipkin server.
publish :: MonadIO m => Zipkin -> m ()
publish :: Zipkin -> m ()
publish Zipkin
z =
  IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Maybe Endpoint -> Tracer -> Request -> Manager -> IO ()
flushSpans (Zipkin -> Maybe Endpoint
zipkinEndpoint Zipkin
z) (Zipkin -> Tracer
zipkinTracer Zipkin
z) (Zipkin -> Request
zipkinRequest Zipkin
z) (Zipkin -> Manager
zipkinManager Zipkin
z)

-- | Convenience method to start a 'Zipkin', run an action, and publish all spans before returning.
with :: MonadUnliftIO m => Settings -> (Zipkin -> m a) -> m a
with :: Settings -> (Zipkin -> m a) -> m a
with Settings
settings Zipkin -> m a
f = do
  Zipkin
zipkin <- Settings -> m Zipkin
forall (m :: * -> *). MonadIO m => Settings -> m Zipkin
new Settings
settings
  Zipkin -> m a
f Zipkin
zipkin m a -> m () -> m a
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`finally` Zipkin -> m ()
forall (m :: * -> *). MonadIO m => Zipkin -> m ()
publish Zipkin
zipkin

-- | Adds a tag to the active span.
tag :: MonadTrace m => Text -> Text -> m ()
tag :: Text -> Text -> m ()
tag Text
key Text
val = Text -> Value -> m ()
forall (m :: * -> *). MonadTrace m => Text -> Value -> m ()
addSpanEntry (Text
publicKeyPrefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
key) (Text -> Value
tagTextValue Text
val)

-- | Adds a tag to a builder. This is a convenience method to use with 'childSpanWith', for example:
--
-- > childSpanWith (addTag "key" "value") "run" $ action
--
-- Note that there is no difference with adding the tag after the span. So the above code is
-- equivalent to:
--
-- > childSpan "run" $ tag "key" "value" >> action
addTag :: Text -> Text -> Builder -> Builder
addTag :: Text -> Text -> Builder -> Builder
addTag Text
key Text
val Builder
bldr =
  Builder
bldr { builderTags :: Map Text Value
builderTags = Text -> Value -> Map Text Value -> Map Text Value
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (Text
publicKeyPrefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
key) (Text -> Value
forall a. ToJSON a => a -> Value
JSON.toJSON Text
val) (Builder -> Map Text Value
builderTags Builder
bldr) }

-- | Adds a producer kind tag to a builder. This is a convenience method to use with 'rootSpanWith',
-- for example:
--
-- > rootSpanWith addProducerKind alwaysSampled "root" $ action
--
-- Use this method if you want to create a root producer span. Otherwise use 'producerSpanWith' to
-- create a sub span with producer kind.
addProducerKind :: Builder -> Builder
addProducerKind :: Builder -> Builder
addProducerKind = Text -> Text -> Builder -> Builder
addTag Text
kindKey Text
producerKindValue

-- | Adds an inherited tag to a builder. Unlike a tag added via 'addTag', this tag:
--
-- * will be inherited by all the span's /local/ children.
-- * can only be added at span construction time.
--
-- For example, to add an ID tag to all spans inside a trace:
--
-- > rootSpanWith (addInheritedTag "id" "abcd-efg") alwaysSampled "run" $ action
addInheritedTag :: Text -> Text -> Builder -> Builder
addInheritedTag :: Text -> Text -> Builder -> Builder
addInheritedTag Text
key Text
val Builder
bldr =
  let bgs :: Map Text Method
bgs = Builder -> Map Text Method
builderBaggages Builder
bldr
  in Builder
bldr { builderBaggages :: Map Text Method
builderBaggages = Text -> Method -> Map Text Method -> Map Text Method
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
key (Text -> Method
T.encodeUtf8 Text
val) Map Text Method
bgs }

-- | Annotates the active span using the current time.
annotate :: MonadTrace m => Text -> m ()
annotate :: Text -> m ()
annotate Text
val = Text -> Value -> m ()
forall (m :: * -> *). MonadTrace m => Text -> Value -> m ()
addSpanEntry Text
"" (Text -> Value
forall a. ToJSON a => a -> Value
logValue Text
val)

-- | Annotates the active span at the given time.
annotateAt :: MonadTrace m => POSIXTime -> Text -> m ()
annotateAt :: NominalDiffTime -> Text -> m ()
annotateAt NominalDiffTime
time Text
val = Text -> Value -> m ()
forall (m :: * -> *). MonadTrace m => Text -> Value -> m ()
addSpanEntry Text
"" (NominalDiffTime -> Text -> Value
forall a. ToJSON a => NominalDiffTime -> a -> Value
logValueAt NominalDiffTime
time Text
val)

-- | Exportable trace information, used for cross-process traces.
data B3 = B3
  { B3 -> TraceID
b3TraceID :: !TraceID
  -- ^ The span's trace ID.
  , B3 -> SpanID
b3SpanID :: !SpanID
  -- ^ The span's ID.
  , B3 -> Bool
b3IsSampled :: !Bool
  -- ^ Whether the span was sampled.
  , B3 -> Bool
b3IsDebug :: !Bool
  -- ^ Whether the span has debug enabled (which implies that the span is sampled).
  , B3 -> Maybe SpanID
b3ParentSpanID :: !(Maybe SpanID)
  -- ^ The span's parent's ID, or 'Nothing' for root spans.
  } deriving (B3 -> B3 -> Bool
(B3 -> B3 -> Bool) -> (B3 -> B3 -> Bool) -> Eq B3
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: B3 -> B3 -> Bool
$c/= :: B3 -> B3 -> Bool
== :: B3 -> B3 -> Bool
$c== :: B3 -> B3 -> Bool
Eq, Eq B3
Eq B3
-> (B3 -> B3 -> Ordering)
-> (B3 -> B3 -> Bool)
-> (B3 -> B3 -> Bool)
-> (B3 -> B3 -> Bool)
-> (B3 -> B3 -> Bool)
-> (B3 -> B3 -> B3)
-> (B3 -> B3 -> B3)
-> Ord B3
B3 -> B3 -> Bool
B3 -> B3 -> Ordering
B3 -> B3 -> B3
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: B3 -> B3 -> B3
$cmin :: B3 -> B3 -> B3
max :: B3 -> B3 -> B3
$cmax :: B3 -> B3 -> B3
>= :: B3 -> B3 -> Bool
$c>= :: B3 -> B3 -> Bool
> :: B3 -> B3 -> Bool
$c> :: B3 -> B3 -> Bool
<= :: B3 -> B3 -> Bool
$c<= :: B3 -> B3 -> Bool
< :: B3 -> B3 -> Bool
$c< :: B3 -> B3 -> Bool
compare :: B3 -> B3 -> Ordering
$ccompare :: B3 -> B3 -> Ordering
$cp1Ord :: Eq B3
Ord, Int -> B3 -> ShowS
[B3] -> ShowS
B3 -> HostName
(Int -> B3 -> ShowS)
-> (B3 -> HostName) -> ([B3] -> ShowS) -> Show B3
forall a.
(Int -> a -> ShowS) -> (a -> HostName) -> ([a] -> ShowS) -> Show a
showList :: [B3] -> ShowS
$cshowList :: [B3] -> ShowS
show :: B3 -> HostName
$cshow :: B3 -> HostName
showsPrec :: Int -> B3 -> ShowS
$cshowsPrec :: Int -> B3 -> ShowS
Show)

traceIDHeader, spanIDHeader, parentSpanIDHeader, sampledHeader, debugHeader :: CI ByteString
traceIDHeader :: HeaderName
traceIDHeader = HeaderName
"X-B3-TraceId"
spanIDHeader :: HeaderName
spanIDHeader = HeaderName
"X-B3-SpanId"
parentSpanIDHeader :: HeaderName
parentSpanIDHeader = HeaderName
"X-B3-ParentSpanId"
sampledHeader :: HeaderName
sampledHeader = HeaderName
"X-B3-Sampled"
debugHeader :: HeaderName
debugHeader = HeaderName
"X-B3-Flags"

-- | Serializes the 'B3' to multiple headers, suitable for HTTP requests. All byte-strings are UTF-8
-- encoded.
b3ToHeaders :: B3 -> Map (CI ByteString) ByteString
b3ToHeaders :: B3 -> Map HeaderName Method
b3ToHeaders (B3 TraceID
traceID SpanID
spanID Bool
isSampled Bool
isDebug Maybe SpanID
mbParentID) =
  let
    defaultKVs :: [(HeaderName, Text)]
defaultKVs = [(HeaderName
traceIDHeader, TraceID -> Text
encodeZipkinTraceID TraceID
traceID), (HeaderName
spanIDHeader, SpanID -> Text
encodeSpanID SpanID
spanID)]
    parentKVs :: [(HeaderName, Text)]
parentKVs = (HeaderName
parentSpanIDHeader,) (Text -> (HeaderName, Text))
-> (SpanID -> Text) -> SpanID -> (HeaderName, Text)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SpanID -> Text
encodeSpanID (SpanID -> (HeaderName, Text)) -> [SpanID] -> [(HeaderName, Text)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe SpanID -> [SpanID]
forall a. Maybe a -> [a]
maybeToList Maybe SpanID
mbParentID
    sampledKVs :: [(HeaderName, Text)]
sampledKVs = case (Bool
isSampled, Bool
isDebug) of
      (Bool
_, Bool
True) -> [(HeaderName
debugHeader, Text
"1")]
      (Bool
True, Bool
_) -> [(HeaderName
sampledHeader, Text
"1")]
      (Bool
False, Bool
_) -> [(HeaderName
sampledHeader, Text
"0")]
  in (Text -> Method) -> Map HeaderName Text -> Map HeaderName Method
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Text -> Method
T.encodeUtf8 (Map HeaderName Text -> Map HeaderName Method)
-> Map HeaderName Text -> Map HeaderName Method
forall a b. (a -> b) -> a -> b
$ [(HeaderName, Text)] -> Map HeaderName Text
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(HeaderName, Text)] -> Map HeaderName Text)
-> [(HeaderName, Text)] -> Map HeaderName Text
forall a b. (a -> b) -> a -> b
$ [(HeaderName, Text)]
defaultKVs [(HeaderName, Text)]
-> [(HeaderName, Text)] -> [(HeaderName, Text)]
forall a. [a] -> [a] -> [a]
++ [(HeaderName, Text)]
parentKVs [(HeaderName, Text)]
-> [(HeaderName, Text)] -> [(HeaderName, Text)]
forall a. [a] -> [a] -> [a]
++ [(HeaderName, Text)]
sampledKVs

-- | Deserializes the 'B3' from multiple headers.
b3FromHeaders :: Map (CI ByteString) ByteString -> Maybe B3
b3FromHeaders :: Map HeaderName Method -> Maybe B3
b3FromHeaders Map HeaderName Method
hdrs = do
  let
    find :: HeaderName -> Maybe Text
find HeaderName
key = Method -> Text
T.decodeUtf8 (Method -> Text) -> Maybe Method -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HeaderName -> Map HeaderName Method -> Maybe Method
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup HeaderName
key Map HeaderName Method
hdrs
    findBool :: Bool -> HeaderName -> Maybe Bool
findBool Bool
def HeaderName
key = case HeaderName -> Maybe Text
find HeaderName
key of
      Maybe Text
Nothing -> Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
def
      Just Text
"1" -> Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
True
      Just Text
"0" -> Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
False
      Maybe Text
_ -> Maybe Bool
forall a. Maybe a
Nothing
  Bool
dbg <- Bool -> HeaderName -> Maybe Bool
findBool Bool
False HeaderName
debugHeader
  Bool
sampled <- Bool -> HeaderName -> Maybe Bool
findBool Bool
dbg HeaderName
sampledHeader
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Bool
sampled Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
False Bool -> Bool -> Bool
&& Bool
dbg)
  TraceID -> SpanID -> Bool -> Bool -> Maybe SpanID -> B3
B3
    (TraceID -> SpanID -> Bool -> Bool -> Maybe SpanID -> B3)
-> Maybe TraceID
-> Maybe (SpanID -> Bool -> Bool -> Maybe SpanID -> B3)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (HeaderName -> Maybe Text
find HeaderName
traceIDHeader Maybe Text -> (Text -> Maybe TraceID) -> Maybe TraceID
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Text -> Maybe TraceID
decodeZipkinTraceID)
    Maybe (SpanID -> Bool -> Bool -> Maybe SpanID -> B3)
-> Maybe SpanID -> Maybe (Bool -> Bool -> Maybe SpanID -> B3)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (HeaderName -> Maybe Text
find HeaderName
spanIDHeader Maybe Text -> (Text -> Maybe SpanID) -> Maybe SpanID
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Text -> Maybe SpanID
decodeSpanID)
    Maybe (Bool -> Bool -> Maybe SpanID -> B3)
-> Maybe Bool -> Maybe (Bool -> Maybe SpanID -> B3)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> Maybe Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
sampled
    Maybe (Bool -> Maybe SpanID -> B3)
-> Maybe Bool -> Maybe (Maybe SpanID -> B3)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> Maybe Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
dbg
    Maybe (Maybe SpanID -> B3) -> Maybe (Maybe SpanID) -> Maybe B3
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe (Maybe SpanID)
-> (Text -> Maybe (Maybe SpanID))
-> Maybe Text
-> Maybe (Maybe SpanID)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Maybe SpanID -> Maybe (Maybe SpanID)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe SpanID
forall a. Maybe a
Nothing) (Maybe SpanID -> Maybe (Maybe SpanID)
forall a. a -> Maybe a
Just (Maybe SpanID -> Maybe (Maybe SpanID))
-> (Text -> Maybe SpanID) -> Text -> Maybe (Maybe SpanID)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Maybe SpanID
decodeSpanID) (HeaderName -> Maybe Text
find HeaderName
parentSpanIDHeader)

-- | Serializes the 'B3' to a single UTF-8 encoded header value. It will typically be set as
-- <https://github.com/apache/incubator-zipkin-b3-propagation#single-header b3 header>.
b3ToHeaderValue :: B3 -> ByteString
b3ToHeaderValue :: B3 -> Method
b3ToHeaderValue (B3 TraceID
traceID SpanID
spanID Bool
isSampled Bool
isDebug Maybe SpanID
mbParentID) =
  let
    state :: Text
state = case (Bool
isSampled, Bool
isDebug) of
      (Bool
_ , Bool
True) -> Text
"d"
      (Bool
True, Bool
_) -> Text
"1"
      (Bool
False, Bool
_) -> Text
"0"
    required :: [Text]
required = [TraceID -> Text
encodeZipkinTraceID TraceID
traceID, SpanID -> Text
encodeSpanID SpanID
spanID, Text
state]
    optional :: [Text]
optional = SpanID -> Text
encodeSpanID (SpanID -> Text) -> [SpanID] -> [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe SpanID -> [SpanID]
forall a. Maybe a -> [a]
maybeToList Maybe SpanID
mbParentID
  in Method -> [Method] -> Method
BS.intercalate Method
"-" ([Method] -> Method) -> [Method] -> Method
forall a b. (a -> b) -> a -> b
$ (Text -> Method) -> [Text] -> [Method]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Text -> Method
T.encodeUtf8 ([Text] -> [Method]) -> [Text] -> [Method]
forall a b. (a -> b) -> a -> b
$ [Text]
required [Text] -> [Text] -> [Text]
forall a. [a] -> [a] -> [a]
++ [Text]
optional

-- | Prefix used to fill up 128-bit if only a 64-bit trace identifier is given.
shortTraceIDPrefix :: Text
shortTraceIDPrefix :: Text
shortTraceIDPrefix = Text
"0000000000000000"

-- | Decodes a zipkin trace ID from a hex-encoded string, returning nothing if it is invalid. Takes
-- into account that the provided string could be a 16 or 32 lower-hex character trace ID. If the
-- given string consists of 16 lower-hex characters 'shortTraceIDPrefix' is used to fil up the
-- 128-bit trace identifier of 'TraceID'.
decodeZipkinTraceID :: Text -> Maybe TraceID
decodeZipkinTraceID :: Text -> Maybe TraceID
decodeZipkinTraceID Text
txt =
  let normalized :: Text
normalized = if Text -> Int
T.length Text
txt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
16 then Text
shortTraceIDPrefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
txt else Text
txt
  in Text -> Maybe TraceID
decodeTraceID Text
normalized

-- | Hex-encodes a trace ID, providing a 16 or 32 lower-hex character zipkin trace ID. A 16
-- lower-hex character string is returned if the first 64-bits of the 'TraceID' are zeros.
encodeZipkinTraceID :: TraceID -> Text
encodeZipkinTraceID :: TraceID -> Text
encodeZipkinTraceID TraceID
traceID =
  let txt :: Text
txt = TraceID -> Text
encodeTraceID TraceID
traceID
  in Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
txt (Maybe Text -> Text) -> Maybe Text -> Text
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Maybe Text
T.stripPrefix Text
shortTraceIDPrefix Text
txt

-- | Deserializes a single header value into a 'B3'.
b3FromHeaderValue :: ByteString -> Maybe B3
b3FromHeaderValue :: Method -> Maybe B3
b3FromHeaderValue Method
bs = case Text -> Text -> [Text]
T.splitOn Text
"-" (Text -> [Text]) -> Text -> [Text]
forall a b. (a -> b) -> a -> b
$ Method -> Text
T.decodeUtf8 Method
bs of
  (Text
traceIDstr:Text
spanIDstr:[Text]
strs) -> do
    TraceID
traceID <- Text -> Maybe TraceID
decodeZipkinTraceID Text
traceIDstr
    SpanID
spanID <- Text -> Maybe SpanID
decodeSpanID Text
spanIDstr
    let buildB3 :: Bool -> Bool -> Maybe SpanID -> B3
buildB3 = TraceID -> SpanID -> Bool -> Bool -> Maybe SpanID -> B3
B3 TraceID
traceID SpanID
spanID
    case [Text]
strs of
      [] -> B3 -> Maybe B3
forall (f :: * -> *) a. Applicative f => a -> f a
pure (B3 -> Maybe B3) -> B3 -> Maybe B3
forall a b. (a -> b) -> a -> b
$ Bool -> Bool -> Maybe SpanID -> B3
buildB3 Bool
False Bool
False Maybe SpanID
forall a. Maybe a
Nothing
      (Text
state:[Text]
strs') -> do
        Maybe SpanID -> B3
buildB3' <- case Text
state of
          Text
"0" -> (Maybe SpanID -> B3) -> Maybe (Maybe SpanID -> B3)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Maybe SpanID -> B3) -> Maybe (Maybe SpanID -> B3))
-> (Maybe SpanID -> B3) -> Maybe (Maybe SpanID -> B3)
forall a b. (a -> b) -> a -> b
$ Bool -> Bool -> Maybe SpanID -> B3
buildB3 Bool
False Bool
False
          Text
"1" -> (Maybe SpanID -> B3) -> Maybe (Maybe SpanID -> B3)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Maybe SpanID -> B3) -> Maybe (Maybe SpanID -> B3))
-> (Maybe SpanID -> B3) -> Maybe (Maybe SpanID -> B3)
forall a b. (a -> b) -> a -> b
$ Bool -> Bool -> Maybe SpanID -> B3
buildB3 Bool
True Bool
False
          Text
"d" -> (Maybe SpanID -> B3) -> Maybe (Maybe SpanID -> B3)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Maybe SpanID -> B3) -> Maybe (Maybe SpanID -> B3))
-> (Maybe SpanID -> B3) -> Maybe (Maybe SpanID -> B3)
forall a b. (a -> b) -> a -> b
$ Bool -> Bool -> Maybe SpanID -> B3
buildB3 Bool
True Bool
True
          Text
_ -> Maybe (Maybe SpanID -> B3)
forall a. Maybe a
Nothing
        case [Text]
strs' of
          [] -> B3 -> Maybe B3
forall (f :: * -> *) a. Applicative f => a -> f a
pure (B3 -> Maybe B3) -> B3 -> Maybe B3
forall a b. (a -> b) -> a -> b
$ Maybe SpanID -> B3
buildB3' Maybe SpanID
forall a. Maybe a
Nothing
          [Text
str] -> Maybe SpanID -> B3
buildB3' (Maybe SpanID -> B3) -> (SpanID -> Maybe SpanID) -> SpanID -> B3
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SpanID -> Maybe SpanID
forall a. a -> Maybe a
Just (SpanID -> B3) -> Maybe SpanID -> Maybe B3
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Maybe SpanID
decodeSpanID Text
str
          [Text]
_ -> Maybe B3
forall a. Maybe a
Nothing
  [Text]
_ -> Maybe B3
forall a. Maybe a
Nothing

b3FromSpan :: Span -> B3
b3FromSpan :: Span -> B3
b3FromSpan Span
s =
  let
    ctx :: Context
ctx = Span -> Context
spanContext Span
s
    refs :: Set Reference
refs = Span -> Set Reference
spanReferences Span
s
  in TraceID -> SpanID -> Bool -> Bool -> Maybe SpanID -> B3
B3 (Context -> TraceID
contextTraceID Context
ctx) (Context -> SpanID
contextSpanID Context
ctx) (Span -> Bool
spanIsSampled Span
s) (Span -> Bool
spanIsDebug Span
s) (Set Reference -> Maybe SpanID
parentID Set Reference
refs)

-- Builder endos

insertTag :: JSON.ToJSON a => Key -> a -> Endo Builder
insertTag :: Text -> a -> Endo Builder
insertTag Text
key a
val =
  (Builder -> Builder) -> Endo Builder
forall a. (a -> a) -> Endo a
Endo ((Builder -> Builder) -> Endo Builder)
-> (Builder -> Builder) -> Endo Builder
forall a b. (a -> b) -> a -> b
$ \Builder
bldr -> Builder
bldr { builderTags :: Map Text Value
builderTags = Text -> Value -> Map Text Value -> Map Text Value
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
key (a -> Value
forall a. ToJSON a => a -> Value
JSON.toJSON a
val) (Builder -> Map Text Value
builderTags Builder
bldr) }

importB3 :: B3 -> Endo Builder
importB3 :: B3 -> Endo Builder
importB3 B3
b3 =
  let
    policy :: SamplingPolicy
policy = if B3 -> Bool
b3IsDebug B3
b3
      then SamplingPolicy
debugEnabled
      else Bool -> SamplingPolicy
sampledWhen (Bool -> SamplingPolicy) -> Bool -> SamplingPolicy
forall a b. (a -> b) -> a -> b
$ B3 -> Bool
b3IsSampled B3
b3
  in (Builder -> Builder) -> Endo Builder
forall a. (a -> a) -> Endo a
Endo ((Builder -> Builder) -> Endo Builder)
-> (Builder -> Builder) -> Endo Builder
forall a b. (a -> b) -> a -> b
$ \Builder
bldr -> Builder
bldr
    { builderTraceID :: Maybe TraceID
builderTraceID = TraceID -> Maybe TraceID
forall a. a -> Maybe a
Just (B3 -> TraceID
b3TraceID B3
b3)
    , builderSamplingPolicy :: Maybe SamplingPolicy
builderSamplingPolicy = SamplingPolicy -> Maybe SamplingPolicy
forall a. a -> Maybe a
Just SamplingPolicy
policy }

-- Prefix added to all user tags. This protects against collisions with internal tags.
publicKeyPrefix :: Text
publicKeyPrefix :: Text
publicKeyPrefix = Text
"Z."

-- Remote endpoint tag key.
endpointKey :: Key
endpointKey :: Text
endpointKey = Text
"z.e"

-- Kind tag key.
kindKey :: Key
kindKey :: Text
kindKey = Text
"z.k"

-- Value that indicates a producer span kind.
producerKindValue :: Text
producerKindValue :: Text
producerKindValue = Text
"PRODUCER"

outgoingSpan :: MonadTrace m => Text -> Endo Builder -> Name -> (Maybe B3 -> m a) -> m a
outgoingSpan :: Text -> Endo Builder -> Text -> (Maybe B3 -> m a) -> m a
outgoingSpan Text
kind Endo Builder
endo Text
name Maybe B3 -> m a
f = (Builder -> Builder) -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
(Builder -> Builder) -> Text -> m a -> m a
childSpanWith (Endo Builder -> Builder -> Builder
forall a. Endo a -> a -> a
appEndo Endo Builder
endo') Text
name m a
actn where
  endo' :: Endo Builder
endo' = Text -> Text -> Endo Builder
forall a. ToJSON a => Text -> a -> Endo Builder
insertTag Text
kindKey Text
kind Endo Builder -> Endo Builder -> Endo Builder
forall a. Semigroup a => a -> a -> a
<> Endo Builder
endo
  actn :: m a
actn = m (Maybe Span)
forall (m :: * -> *). MonadTrace m => m (Maybe Span)
activeSpan m (Maybe Span) -> (Maybe Span -> m a) -> m a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe Span
Nothing -> Maybe B3 -> m a
f Maybe B3
forall a. Maybe a
Nothing
    Just Span
spn -> Maybe B3 -> m a
f (Maybe B3 -> m a) -> Maybe B3 -> m a
forall a b. (a -> b) -> a -> b
$ B3 -> Maybe B3
forall a. a -> Maybe a
Just (B3 -> Maybe B3) -> B3 -> Maybe B3
forall a b. (a -> b) -> a -> b
$ Span -> B3
b3FromSpan Span
spn

-- | Generates a child span with @CLIENT@ kind. This function also provides the corresponding 'B3'
-- (or 'Nothing' if tracing is inactive) so that it can be forwarded to the server. For example, to
-- emit an HTTP request and forward the trace information in the headers:
--
-- > import Network.HTTP.Simple
-- >
-- > clientSpan "api-call" $ \(Just b3) -> $ do
-- >   res <- httpBS "http://host/api" & addRequestHeader "b3" (b3ToHeaderValue b3)
-- >   process res -- Do something with the response.
clientSpan :: MonadTrace m => Name -> (Maybe B3 -> m a) -> m a
clientSpan :: Text -> (Maybe B3 -> m a) -> m a
clientSpan = (Builder -> Builder) -> Text -> (Maybe B3 -> m a) -> m a
forall (m :: * -> *) a.
MonadTrace m =>
(Builder -> Builder) -> Text -> (Maybe B3 -> m a) -> m a
clientSpanWith Builder -> Builder
forall a. a -> a
id

-- | Generates a client span, optionally modifying the span's builder. This can be useful in
-- combination with 'addEndpoint' if the remote server does not have tracing enabled.
clientSpanWith :: MonadTrace m => (Builder -> Builder) -> Name -> (Maybe B3 -> m a) -> m a
clientSpanWith :: (Builder -> Builder) -> Text -> (Maybe B3 -> m a) -> m a
clientSpanWith Builder -> Builder
f = Text -> Endo Builder -> Text -> (Maybe B3 -> m a) -> m a
forall (m :: * -> *) a.
MonadTrace m =>
Text -> Endo Builder -> Text -> (Maybe B3 -> m a) -> m a
outgoingSpan Text
"CLIENT" ((Builder -> Builder) -> Endo Builder
forall a. (a -> a) -> Endo a
Endo Builder -> Builder
f)

-- | Generates a child span with @PRODUCER@ kind. This function also provides the corresponding 'B3'
-- so that it can be forwarded to the consumer.
producerSpanWith :: MonadTrace m => (Builder -> Builder) -> Name -> (Maybe B3 -> m a) -> m a
producerSpanWith :: (Builder -> Builder) -> Text -> (Maybe B3 -> m a) -> m a
producerSpanWith Builder -> Builder
f = Text -> Endo Builder -> Text -> (Maybe B3 -> m a) -> m a
forall (m :: * -> *) a.
MonadTrace m =>
Text -> Endo Builder -> Text -> (Maybe B3 -> m a) -> m a
outgoingSpan Text
producerKindValue ((Builder -> Builder) -> Endo Builder
forall a. (a -> a) -> Endo a
Endo Builder -> Builder
f)

incomingSpan :: MonadTrace m => Text -> B3 -> Endo Builder -> m a -> m a
incomingSpan :: Text -> B3 -> Endo Builder -> m a -> m a
incomingSpan Text
kind B3
b3 Endo Builder
endo m a
actn =
  let bldr :: Builder
bldr = Endo Builder -> Builder -> Builder
forall a. Endo a -> a -> a
appEndo (Text -> Text -> Endo Builder
forall a. ToJSON a => Text -> a -> Endo Builder
insertTag Text
kindKey Text
kind Endo Builder -> Endo Builder -> Endo Builder
forall a. Semigroup a => a -> a -> a
<> B3 -> Endo Builder
importB3 B3
b3 Endo Builder -> Endo Builder -> Endo Builder
forall a. Semigroup a => a -> a -> a
<> Endo Builder
endo) (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ Text -> Builder
builder Text
""
  in Builder -> m a -> m a
forall (m :: * -> *) a. MonadTrace m => Builder -> m a -> m a
trace Builder
bldr m a
actn

-- | Generates a child span with @SERVER@ kind. The client's 'B3' should be provided as input,
-- for example parsed using 'b3FromHeaders'.
serverSpan :: MonadTrace m => B3 -> m a -> m a
serverSpan :: B3 -> m a -> m a
serverSpan = (Builder -> Builder) -> B3 -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
(Builder -> Builder) -> B3 -> m a -> m a
serverSpanWith Builder -> Builder
forall a. a -> a
id

-- | Generates a child span with @SERVER@ kind, optionally modifying the span's builder. This can
-- be useful in combination with 'addEndpoint' if the remote client does not have tracing enabled.
-- The clients's 'B3' should be provided as input. Client and server annotations go on the same
-- span - it means that they share their span ID.
serverSpanWith :: MonadTrace m => (Builder -> Builder) -> B3 -> m a -> m a
serverSpanWith :: (Builder -> Builder) -> B3 -> m a -> m a
serverSpanWith Builder -> Builder
f B3
b3 =
  let endo :: Endo Builder
endo = (Builder -> Builder) -> Endo Builder
forall a. (a -> a) -> Endo a
Endo ((Builder -> Builder) -> Endo Builder)
-> (Builder -> Builder) -> Endo Builder
forall a b. (a -> b) -> a -> b
$ \Builder
bldr -> Builder -> Builder
f (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ Builder
bldr { builderSpanID :: Maybe SpanID
builderSpanID = SpanID -> Maybe SpanID
forall a. a -> Maybe a
Just (B3 -> SpanID
b3SpanID B3
b3) }
  in Text -> B3 -> Endo Builder -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
Text -> B3 -> Endo Builder -> m a -> m a
incomingSpan Text
"SERVER" B3
b3 Endo Builder
endo

-- | Generates a child span with @CONSUMER@ kind, optionally modifying the span's builder. The
-- producer's 'B3' should be provided as input. The generated span will have its parent ID set to
-- the input B3's span ID.
consumerSpanWith :: MonadTrace m => (Builder -> Builder) -> B3 -> m a -> m a
consumerSpanWith :: (Builder -> Builder) -> B3 -> m a -> m a
consumerSpanWith Builder -> Builder
f B3
b3 =
  let endo :: Endo Builder
endo = (Builder -> Builder) -> Endo Builder
forall a. (a -> a) -> Endo a
Endo ((Builder -> Builder) -> Endo Builder)
-> (Builder -> Builder) -> Endo Builder
forall a b. (a -> b) -> a -> b
$ \Builder
bldr -> Builder -> Builder
f (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ Builder
bldr { builderReferences :: Set Reference
builderReferences = Reference -> Set Reference
forall a. a -> Set a
Set.singleton (SpanID -> Reference
ChildOf (SpanID -> Reference) -> SpanID -> Reference
forall a b. (a -> b) -> a -> b
$ B3 -> SpanID
b3SpanID B3
b3) }
  in Text -> B3 -> Endo Builder -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
Text -> B3 -> Endo Builder -> m a -> m a
incomingSpan Text
"CONSUMER" B3
b3 Endo Builder
endo

-- | Information about a hosted service, included in spans and visible in the Zipkin UI.
data Endpoint = Endpoint
  { Endpoint -> Maybe Text
endpointService :: !(Maybe Text)
  -- ^ The endpoint's service name.
  , Endpoint -> Maybe Int
endpointPort :: !(Maybe Int)
  -- ^ The endpoint's port, if applicable and known.
  , Endpoint -> Maybe Text
endpointIPv4 :: !(Maybe Text)
  -- ^ The endpoint's IPv4 address.
  , Endpoint -> Maybe Text
endpointIPv6 :: !(Maybe Text)
  -- ^ The endpoint's IPv6 address.
  } deriving (Endpoint -> Endpoint -> Bool
(Endpoint -> Endpoint -> Bool)
-> (Endpoint -> Endpoint -> Bool) -> Eq Endpoint
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Endpoint -> Endpoint -> Bool
$c/= :: Endpoint -> Endpoint -> Bool
== :: Endpoint -> Endpoint -> Bool
$c== :: Endpoint -> Endpoint -> Bool
Eq, Eq Endpoint
Eq Endpoint
-> (Endpoint -> Endpoint -> Ordering)
-> (Endpoint -> Endpoint -> Bool)
-> (Endpoint -> Endpoint -> Bool)
-> (Endpoint -> Endpoint -> Bool)
-> (Endpoint -> Endpoint -> Bool)
-> (Endpoint -> Endpoint -> Endpoint)
-> (Endpoint -> Endpoint -> Endpoint)
-> Ord Endpoint
Endpoint -> Endpoint -> Bool
Endpoint -> Endpoint -> Ordering
Endpoint -> Endpoint -> Endpoint
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Endpoint -> Endpoint -> Endpoint
$cmin :: Endpoint -> Endpoint -> Endpoint
max :: Endpoint -> Endpoint -> Endpoint
$cmax :: Endpoint -> Endpoint -> Endpoint
>= :: Endpoint -> Endpoint -> Bool
$c>= :: Endpoint -> Endpoint -> Bool
> :: Endpoint -> Endpoint -> Bool
$c> :: Endpoint -> Endpoint -> Bool
<= :: Endpoint -> Endpoint -> Bool
$c<= :: Endpoint -> Endpoint -> Bool
< :: Endpoint -> Endpoint -> Bool
$c< :: Endpoint -> Endpoint -> Bool
compare :: Endpoint -> Endpoint -> Ordering
$ccompare :: Endpoint -> Endpoint -> Ordering
$cp1Ord :: Eq Endpoint
Ord, Int -> Endpoint -> ShowS
[Endpoint] -> ShowS
Endpoint -> HostName
(Int -> Endpoint -> ShowS)
-> (Endpoint -> HostName) -> ([Endpoint] -> ShowS) -> Show Endpoint
forall a.
(Int -> a -> ShowS) -> (a -> HostName) -> ([a] -> ShowS) -> Show a
showList :: [Endpoint] -> ShowS
$cshowList :: [Endpoint] -> ShowS
show :: Endpoint -> HostName
$cshow :: Endpoint -> HostName
showsPrec :: Int -> Endpoint -> ShowS
$cshowsPrec :: Int -> Endpoint -> ShowS
Show)

-- | An empty endpoint.
defaultEndpoint :: Endpoint
defaultEndpoint :: Endpoint
defaultEndpoint = Maybe Text -> Maybe Int -> Maybe Text -> Maybe Text -> Endpoint
Endpoint Maybe Text
forall a. Maybe a
Nothing Maybe Int
forall a. Maybe a
Nothing Maybe Text
forall a. Maybe a
Nothing Maybe Text
forall a. Maybe a
Nothing

-- | Adds a remote endpoint to a builder. This is mostly useful when generating cross-process spans
-- where the remote endpoint is not already traced (otherwise Zipkin will associate the spans
-- correctly automatically). For example when emitting a request to an outside server:
--
-- > clientSpanWith (addEndpoint "outside-api") -- ...
addEndpoint :: Endpoint -> Builder -> Builder
addEndpoint :: Endpoint -> Builder -> Builder
addEndpoint = Endo Builder -> Builder -> Builder
forall a. Endo a -> a -> a
appEndo (Endo Builder -> Builder -> Builder)
-> (Endpoint -> Endo Builder) -> Endpoint -> Builder -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Endpoint -> Endo Builder
forall a. ToJSON a => Text -> a -> Endo Builder
insertTag Text
endpointKey

-- | Generates an endpoint with the given string as service.
instance IsString Endpoint where
  fromString :: HostName -> Endpoint
fromString HostName
s = Endpoint
defaultEndpoint { endpointService :: Maybe Text
endpointService = Text -> Maybe Text
forall a. a -> Maybe a
Just (HostName -> Text
T.pack HostName
s) }

instance JSON.ToJSON Endpoint where
  toJSON :: Endpoint -> Value
toJSON (Endpoint Maybe Text
mbSvc Maybe Int
mbPort Maybe Text
mbIPv4 Maybe Text
mbIPv6) = [Pair] -> Value
JSON.object ([Pair] -> Value) -> [Pair] -> Value
forall a b. (a -> b) -> a -> b
$ [Maybe Pair] -> [Pair]
forall a. [Maybe a] -> [a]
catMaybes
    [ (Text
"serviceName" Text -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..=) (Text -> Pair) -> Maybe Text -> Maybe Pair
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Text
mbSvc
    , (Text
"port" Text -> Int -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..=) (Int -> Pair) -> Maybe Int -> Maybe Pair
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Int
mbPort
    , (Text
"ipv4" Text -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..=) (Text -> Pair) -> Maybe Text -> Maybe Pair
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Text
mbIPv4
    , (Text
"ipv6" Text -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..=) (Text -> Pair) -> Maybe Text -> Maybe Pair
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Text
mbIPv6 ]

parentID :: Set Reference -> Maybe SpanID
parentID :: Set Reference -> Maybe SpanID
parentID = [SpanID] -> Maybe SpanID
forall a. [a] -> Maybe a
listToMaybe ([SpanID] -> Maybe SpanID)
-> (Set Reference -> [SpanID]) -> Set Reference -> Maybe SpanID
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe SpanID] -> [SpanID]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe SpanID] -> [SpanID])
-> (Set Reference -> [Maybe SpanID]) -> Set Reference -> [SpanID]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Reference -> Maybe SpanID) -> [Reference] -> [Maybe SpanID]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Reference -> Maybe SpanID
go ([Reference] -> [Maybe SpanID])
-> (Set Reference -> [Reference])
-> Set Reference
-> [Maybe SpanID]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set Reference -> [Reference]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList where
  go :: Reference -> Maybe SpanID
go (ChildOf SpanID
d) = SpanID -> Maybe SpanID
forall a. a -> Maybe a
Just SpanID
d
  go Reference
_ = Maybe SpanID
forall a. Maybe a
Nothing

data ZipkinAnnotation = ZipkinAnnotation !POSIXTime !JSON.Value

instance JSON.ToJSON ZipkinAnnotation where
  toJSON :: ZipkinAnnotation -> Value
toJSON (ZipkinAnnotation NominalDiffTime
t Value
v) = [Pair] -> Value
JSON.object
    [ Text
"timestamp" Text -> Int64 -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..= NominalDiffTime -> Int64
forall a. Integral a => NominalDiffTime -> a
microSeconds @Int64 NominalDiffTime
t
    , Text
"value" Text -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..= Value
v ]

-- Internal type used to encode spans in the <https://zipkin.apache.org/zipkin-api/#/ format>
-- expected by Zipkin.
data ZipkinSpan = ZipkinSpan !(Maybe Endpoint) !Sample

publicTags :: Tags -> Map Text JSON.Value
publicTags :: Map Text Value -> Map Text Value
publicTags = [Pair] -> Map Text Value
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([Pair] -> Map Text Value)
-> (Map Text Value -> [Pair]) -> Map Text Value -> Map Text Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe Pair] -> [Pair]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe Pair] -> [Pair])
-> (Map Text Value -> [Maybe Pair]) -> Map Text Value -> [Pair]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Pair -> Maybe Pair) -> [Pair] -> [Maybe Pair]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Pair -> Maybe Pair
forall b. (Text, b) -> Maybe (Text, b)
go ([Pair] -> [Maybe Pair])
-> (Map Text Value -> [Pair]) -> Map Text Value -> [Maybe Pair]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Text Value -> [Pair]
forall k a. Map k a -> [(k, a)]
Map.assocs where
  go :: (Text, b) -> Maybe (Text, b)
go (Text
k, b
v) = case Text -> Text -> Maybe Text
T.stripPrefix Text
publicKeyPrefix Text
k of
    Maybe Text
Nothing -> Maybe (Text, b)
forall a. Maybe a
Nothing
    Just Text
k' -> (Text, b) -> Maybe (Text, b)
forall a. a -> Maybe a
Just (Text
k', b
v)

instance JSON.ToJSON ZipkinSpan where
  toJSON :: ZipkinSpan -> Value
toJSON (ZipkinSpan Maybe Endpoint
mbEpt (Sample Span
spn Map Text Value
tags Logs
logs NominalDiffTime
start NominalDiffTime
duration)) =
    let
      ctx :: Context
ctx = Span -> Context
spanContext Span
spn
      requiredKVs :: [Pair]
requiredKVs =
        [ Text
"traceId" Text -> TraceID -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..= Context -> TraceID
contextTraceID Context
ctx
        , Text
"name" Text -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..= Span -> Text
spanName Span
spn
        , Text
"id" Text -> SpanID -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..= Context -> SpanID
contextSpanID Context
ctx
        , Text
"timestamp" Text -> Int64 -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..= NominalDiffTime -> Int64
forall a. Integral a => NominalDiffTime -> a
microSeconds @Int64 NominalDiffTime
start
        , Text
"duration" Text -> Int64 -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..= NominalDiffTime -> Int64
forall a. Integral a => NominalDiffTime -> a
microSeconds @Int64 NominalDiffTime
duration
        , Text
"debug" Text -> Bool -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..= Span -> Bool
spanIsDebug Span
spn
        , Text
"tags" Text -> Map Text Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..= (Map Text Value -> Map Text Value
publicTags Map Text Value
tags Map Text Value -> Map Text Value -> Map Text Value
forall a. Semigroup a => a -> a -> a
<> (Text -> Value
forall a. ToJSON a => a -> Value
JSON.toJSON (Text -> Value) -> (Method -> Text) -> Method -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Method -> Text
T.decodeUtf8 (Method -> Value) -> Map Text Method -> Map Text Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Context -> Map Text Method
contextBaggages Context
ctx))
        , Text
"annotations" Text -> [ZipkinAnnotation] -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..= ((NominalDiffTime, Text, Value) -> ZipkinAnnotation)
-> Logs -> [ZipkinAnnotation]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(NominalDiffTime
t, Text
_, Value
v) -> NominalDiffTime -> Value -> ZipkinAnnotation
ZipkinAnnotation NominalDiffTime
t Value
v) Logs
logs ]
      optionalKVs :: [Pair]
optionalKVs = [Maybe Pair] -> [Pair]
forall a. [Maybe a] -> [a]
catMaybes
        [ (Text
"parentId" Text -> SpanID -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..=) (SpanID -> Pair) -> Maybe SpanID -> Maybe Pair
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Reference -> Maybe SpanID
parentID (Span -> Set Reference
spanReferences Span
spn)
        , (Text
"localEndpoint" Text -> Endpoint -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..=) (Endpoint -> Pair) -> Maybe Endpoint -> Maybe Pair
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Endpoint
mbEpt
        , (Text
"remoteEndpoint" Text -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..=) (Value -> Pair) -> Maybe Value -> Maybe Pair
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Map Text Value -> Maybe Value
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
endpointKey Map Text Value
tags
        , (Text
"kind" Text -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Text -> v -> kv
JSON..=) (Value -> Pair) -> Maybe Value -> Maybe Pair
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Map Text Value -> Maybe Value
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
kindKey Map Text Value
tags ]
    in [Pair] -> Value
JSON.object ([Pair] -> Value) -> [Pair] -> Value
forall a b. (a -> b) -> a -> b
$ [Pair]
requiredKVs [Pair] -> [Pair] -> [Pair]
forall a. [a] -> [a] -> [a]
++ [Pair]
optionalKVs

microSeconds :: Integral a => NominalDiffTime -> a
microSeconds :: NominalDiffTime -> a
microSeconds = NominalDiffTime -> a
forall a b. (RealFrac a, Integral b) => a -> b
round (NominalDiffTime -> a)
-> (NominalDiffTime -> NominalDiffTime) -> NominalDiffTime -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
* NominalDiffTime
1000000)