{-# language FlexibleInstances     #-}
{-# language MultiParamTypeClasses #-}
{-# language OverloadedStrings     #-}
{-# language PolyKinds             #-}
{-# language UndecidableInstances  #-}
{-# language ViewPatterns          #-}
{-|
Description : Distributed tracing for Mu

This module injects distributed tracing
for Mu servers. Currently it only supports
Zipkin as backend.

In order to use this module, you need to
follow these steps:

1. Establish a connection with 'newZipkin'.
2. Wrap the server using 'zipkin', giving
   information for the root.
3. Run the server using the transformer version
   of your protocol, like |grpcAppTrans|.
-}
module Mu.Instrumentation.Tracing (
  -- * Distributed tracing
  MuTracing(..)
, zipkin
, runZipkin
  -- ** Establish connection
, newZipkin
, defaultZipkinSettings
, Settings(..)
  -- * Useful re-exports
, module Monitor.Tracing
) where

import           Control.Applicative       ((<|>))
import           Control.Monad.IO.Class
import           Control.Monad.Trace
import           Control.Monad.Trace.Class
import qualified Data.Map.Strict           as M
import           Data.Text
import           Monitor.Tracing
import           Monitor.Tracing.Zipkin
import           Mu.Rpc
import           Mu.Server

data MuTracing
  = MuTracing {
      MuTracing -> SamplingPolicy
samplingPolicy :: SamplingPolicy
    , MuTracing -> Text
rootName       :: Text
    }

-- | Runs with a given 'Zipkin' connection.
--   You can create one with 'newZipkin'.
runZipkin :: Zipkin -> TraceT m a -> m a
runZipkin :: Zipkin -> TraceT m a -> m a
runZipkin = (TraceT m a -> Zipkin -> m a) -> Zipkin -> TraceT m a -> m a
forall a b c. (a -> b -> c) -> b -> a -> c
flip TraceT m a -> Zipkin -> m a
forall (m :: * -> *) a. TraceT m a -> Zipkin -> m a
run

-- | Create a new connection to 'Zipkin'.
newZipkin :: Settings -> IO Zipkin
newZipkin :: Settings -> IO Zipkin
newZipkin = Settings -> IO Zipkin
forall (m :: * -> *). MonadIO m => Settings -> m Zipkin
new

defaultZipkinSettings :: Settings
defaultZipkinSettings :: Settings
defaultZipkinSettings = Settings
defaultSettings

-- | Wraps a server to do distributed tracing
--   using 'Zipkin' as backend.
zipkin :: (MonadIO m, MonadTrace m)
       => MuTracing -> ServerT chn i p m topHs -> ServerT chn i p m topHs
zipkin :: MuTracing -> ServerT chn i p m topHs -> ServerT chn i p m topHs
zipkin MuTracing
m = (forall a. RpcInfo i -> m a -> m a)
-> ServerT chn i p m topHs -> ServerT chn i p m topHs
forall snm mnm anm (chn :: ServiceChain snm) info
       (p :: Package snm mnm anm (TypeRef snm)) (m :: * -> *)
       (topHs :: [[*]]).
(forall a. RpcInfo info -> m a -> m a)
-> ServerT chn info p m topHs -> ServerT chn info p m topHs
wrapServer (MuTracing -> RpcInfo i -> m a -> m a
forall (m :: * -> *) i a.
(MonadIO m, MonadTrace m) =>
MuTracing -> RpcInfo i -> m a -> m a
zipkinTracing MuTracing
m)

zipkinTracing :: (MonadIO m, MonadTrace m)
              => MuTracing -> RpcInfo i -> m a -> m a
zipkinTracing :: MuTracing -> RpcInfo i -> m a -> m a
zipkinTracing MuTracing
zpk RpcInfo i
NoRpcInfo m a
h =
  SamplingPolicy -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
SamplingPolicy -> Text -> m a -> m a
rootSpan (MuTracing -> SamplingPolicy
samplingPolicy MuTracing
zpk) (MuTracing -> Text
rootName MuTracing
zpk) m a
h
zipkinTracing MuTracing
zpk (RpcInfo Package Text Text Text TyInfo
_ Service Text Text Text TyInfo
_ Method Text Text Text TyInfo
_ (RequestHeaders -> Map HeaderName ByteString
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList -> Map HeaderName ByteString
hdrs) i
_) m a
h =
  case Maybe B3
getB3 of
    Maybe B3
Nothing  -> SamplingPolicy -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
SamplingPolicy -> Text -> m a -> m a
rootSpan (MuTracing -> SamplingPolicy
samplingPolicy MuTracing
zpk) (MuTracing -> Text
rootName MuTracing
zpk) m a
h
    Just B3
spn -> B3 -> m a -> m a
forall (m :: * -> *) a. MonadTrace m => B3 -> m a -> m a
serverSpan B3
spn m a
h
  where getB3 :: Maybe B3
getB3 =   (ByteString -> Maybe B3
b3FromHeaderValue (ByteString -> Maybe B3) -> Maybe ByteString -> Maybe B3
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HeaderName -> Map HeaderName ByteString -> Maybe ByteString
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup HeaderName
"b3" Map HeaderName ByteString
hdrs)
              Maybe B3 -> Maybe B3 -> Maybe B3
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Map HeaderName ByteString -> Maybe B3
b3FromHeaders Map HeaderName ByteString
hdrs