{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Network.HTTP.Client.OpenTracing
    ( httpTraced
    , httpTraced'
    )
where

import           Control.Applicative
import           Control.Lens                 (over, set, view)
import           Control.Monad.IO.Class
import           Control.Monad.Reader
import           Data.Semigroup               ((<>))
import qualified Data.Text                    as Text
import           Data.Text.Encoding           (decodeUtf8)
import           Network.HTTP.Client.Internal
    ( Manager (mModifyRequest, mModifyResponse)
    , Request (..)
    , getUri
    , responseStatus
    )
import           OpenTracing                  hiding (sampled)
import qualified OpenTracing.Propagation      as Propagation
import qualified OpenTracing.Tracer           as Tracer
import           Prelude                      hiding (span)

-- |
--
-- >>> :{
-- mgr <- newManager defaultManagerSettings
-- rq1 <- parseRequest "http://service1.local/foo"
-- rq2 <- parseRequest "http://service2.local/bar"
-- traced (spanOpts "toplevel" mempty) $ \parent -> do
--     rpc1 <- httpTraced (childOf parent) rq1 mgr httpLbs
--     rpc2 <- httpTraced (childOf parent <> followsFrom (tracedSpan rpc1))
--                        rq2 mgr httpLbs
--     return [tracedResult rpc1, tracedResult rpc2]
-- :}
--
httpTraced
    :: ( HasCarrier       Headers p
       , MonadOpenTracing r       p m
       , MonadIO                    m
       )
    => SpanRefs
    -> Request
    -> Manager
    -> (Request -> Manager -> IO a)
    -> m (Traced a)
httpTraced :: SpanRefs
-> Request
-> Manager
-> (Request -> Manager -> IO a)
-> m (Traced a)
httpTraced SpanRefs
refs Request
req Manager
mgr Request -> Manager -> IO a
f = do
    (Tracer
t,Propagation p
p) <- (Tracer -> Propagation p -> (Tracer, Propagation p))
-> m Tracer -> m (Propagation p) -> m (Tracer, Propagation p)
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (,) (Getting Tracer r Tracer -> m Tracer
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Tracer r Tracer
forall a r. HasTracer a => Getting r a Tracer
tracer) (Getting (Propagation p) r (Propagation p) -> m (Propagation p)
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting (Propagation p) r (Propagation p)
forall a (p :: [*]) r.
HasPropagation a p =>
Getting r a (Propagation p)
propagation)
    IO (Traced a) -> m (Traced a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Traced a) -> m (Traced a)) -> IO (Traced a) -> m (Traced a)
forall a b. (a -> b) -> a -> b
$ Tracer
-> Propagation p
-> SpanRefs
-> Request
-> Manager
-> (Request -> Manager -> IO a)
-> IO (Traced a)
forall (p :: [*]) a.
HasCarrier Headers p =>
Tracer
-> Propagation p
-> SpanRefs
-> Request
-> Manager
-> (Request -> Manager -> IO a)
-> IO (Traced a)
httpTraced' Tracer
t Propagation p
p SpanRefs
refs Request
req Manager
mgr Request -> Manager -> IO a
f

httpTraced'
    :: HasCarrier Headers p
    => Tracer
    -> Propagation        p
    -> SpanRefs
    -> Request
    -> Manager
    -> (Request -> Manager -> IO a)
    -> IO (Traced a)
httpTraced' :: Tracer
-> Propagation p
-> SpanRefs
-> Request
-> Manager
-> (Request -> Manager -> IO a)
-> IO (Traced a)
httpTraced' Tracer
t Propagation p
p SpanRefs
refs Request
req Manager
mgr Request -> Manager -> IO a
f = do
    Maybe Sampled
sampled <- (Reference -> Sampled) -> Maybe Reference -> Maybe Sampled
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Getting Sampled SpanContext Sampled -> SpanContext -> Sampled
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Sampled SpanContext Sampled
Lens' SpanContext Sampled
ctxSampled (SpanContext -> Sampled)
-> (Reference -> SpanContext) -> Reference -> Sampled
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Reference -> SpanContext
refCtx) (Maybe Reference -> Maybe Sampled)
-> ([Reference] -> Maybe Reference) -> [Reference] -> Maybe Sampled
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Reference] -> Maybe Reference
forall (t :: * -> *). Foldable t => t Reference -> Maybe Reference
findParent ([Reference] -> Maybe Sampled)
-> IO [Reference] -> IO (Maybe Sampled)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SpanRefs -> IO [Reference]
freezeRefs SpanRefs
refs

    let opt :: SpanOpts
opt = ASetter SpanOpts SpanOpts (Maybe Sampled) (Maybe Sampled)
-> Maybe Sampled -> SpanOpts -> SpanOpts
forall s t a b. ASetter s t a b -> b -> s -> t
set ASetter SpanOpts SpanOpts (Maybe Sampled) (Maybe Sampled)
Lens' SpanOpts (Maybe Sampled)
spanOptSampled Maybe Sampled
sampled
            (SpanOpts -> SpanOpts)
-> (SpanOpts -> SpanOpts) -> SpanOpts -> SpanOpts
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASetter SpanOpts SpanOpts [Tag] [Tag]
-> [Tag] -> SpanOpts -> SpanOpts
forall s t a b. ASetter s t a b -> b -> s -> t
set ASetter SpanOpts SpanOpts [Tag] [Tag]
Lens' SpanOpts [Tag]
spanOptTags
                  [ Method -> Tag
HttpMethod  (Request -> Method
method Request
req)
                  , Text -> Tag
HttpUrl     (String -> Text
Text.pack (String -> Text) -> (URI -> String) -> URI -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. URI -> String
forall a. Show a => a -> String
show (URI -> Text) -> URI -> Text
forall a b. (a -> b) -> a -> b
$ Request -> URI
getUri Request
req)
                  , Text -> Tag
PeerAddress (Method -> Text
decodeUtf8 (Request -> Method
host Request
req))
                  , SpanKinds -> Tag
SpanKind    SpanKinds
RPCClient
                  ]
            (SpanOpts -> SpanOpts) -> SpanOpts -> SpanOpts
forall a b. (a -> b) -> a -> b
$ Text -> SpanRefs -> SpanOpts
spanOpts (Method -> Text
decodeUtf8 (Request -> Method
path Request
req)) SpanRefs
refs

    Tracer -> SpanOpts -> (ActiveSpan -> IO a) -> IO (Traced a)
forall t (m :: * -> *) a.
(HasTracer t, MonadMask m, MonadIO m) =>
t -> SpanOpts -> (ActiveSpan -> m a) -> m (Traced a)
Tracer.traced Tracer
t SpanOpts
opt ((ActiveSpan -> IO a) -> IO (Traced a))
-> (ActiveSpan -> IO a) -> IO (Traced a)
forall a b. (a -> b) -> a -> b
$ \ActiveSpan
span ->
        let mgr' :: Manager
mgr' = ActiveSpan -> Manager
modMgr ActiveSpan
span
         in Request -> Manager -> IO a
f Request
req { requestManagerOverride :: Maybe Manager
requestManagerOverride = Manager -> Maybe Manager
forall a. a -> Maybe a
Just Manager
mgr' } Manager
mgr'
  where
    modMgr :: ActiveSpan -> Manager
modMgr ActiveSpan
span = Manager
mgr
        { mModifyRequest :: Request -> IO Request
mModifyRequest  = \Request
rq ->
            Request -> SpanContext -> Request
inj Request
rq (SpanContext -> Request)
-> (Span -> SpanContext) -> Span -> Request
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting SpanContext Span SpanContext -> Span -> SpanContext
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting SpanContext Span SpanContext
forall a. HasSpanFields a => Lens' a SpanContext
spanContext (Span -> Request) -> IO Span -> IO Request
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ActiveSpan -> IO Span
forall (m :: * -> *). MonadIO m => ActiveSpan -> m Span
readActiveSpan ActiveSpan
span

        , mModifyResponse :: Response BodyReader -> IO (Response BodyReader)
mModifyResponse = \Response BodyReader
rs -> do
            ActiveSpan -> (Span -> Span) -> IO ()
forall (m :: * -> *).
MonadIO m =>
ActiveSpan -> (Span -> Span) -> m ()
modifyActiveSpan ActiveSpan
span ((Span -> Span) -> IO ()) -> (Span -> Span) -> IO ()
forall a b. (a -> b) -> a -> b
$
                ASetter Span Span Tags Tags -> (Tags -> Tags) -> Span -> Span
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter Span Span Tags Tags
forall a. HasSpanFields a => Lens' a Tags
spanTags (Tag -> Tags -> Tags
setTag (Status -> Tag
HttpStatusCode (Response BodyReader -> Status
forall body. Response body -> Status
responseStatus Response BodyReader
rs)))
            Response BodyReader -> IO (Response BodyReader)
forall (m :: * -> *) a. Monad m => a -> m a
return Response BodyReader
rs
        }

    inj :: Request -> SpanContext -> Request
inj Request
rq SpanContext
ctx = Request
rq
        { requestHeaders :: Headers
requestHeaders = Request -> Headers
requestHeaders Request
rq Headers -> Headers -> Headers
forall a. Semigroup a => a -> a -> a
<> Propagation p -> SpanContext -> Headers
forall c r (p :: [*]).
(HasCarrier c p, HasPropagation r p) =>
r -> SpanContext -> c
Propagation.inject Propagation p
p SpanContext
ctx
        }