{-# 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 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 :: forall (p :: [*]) r (m :: * -> *) a.
(HasCarrier Headers p, MonadOpenTracing r p m, MonadIO m) =>
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) <- forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (,) (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall a r. HasTracer a => Getting r a Tracer
tracer) (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall a (p :: [*]) r.
HasPropagation a p =>
Getting r a (Propagation p)
propagation)
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ 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' :: 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 = do
    Maybe Sampled
sampled <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' SpanContext Sampled
ctxSampled forall b c a. (b -> c) -> (a -> b) -> a -> c
. Reference -> SpanContext
refCtx) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *). Foldable t => t Reference -> Maybe Reference
findParent forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SpanRefs -> IO [Reference]
freezeRefs SpanRefs
refs

    let opt :: SpanOpts
opt = forall s t a b. ASetter s t a b -> b -> s -> t
set Lens' SpanOpts (Maybe Sampled)
spanOptSampled Maybe Sampled
sampled
            forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s t a b. ASetter s t a b -> b -> s -> t
set Lens' SpanOpts [Tag]
spanOptTags
                  [ Method -> Tag
HttpMethod  (Request -> Method
method Request
req)
                  , Text -> Tag
HttpUrl     (String -> Text
Text.pack forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show 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
                  ]
            forall a b. (a -> b) -> a -> b
$ Text -> SpanRefs -> SpanOpts
spanOpts (Method -> Text
decodeUtf8 (Request -> Method
path Request
req)) SpanRefs
refs

    forall t (m :: * -> *) a.
(HasTracer t, MonadMask m, MonadIO m) =>
t -> SpanOpts -> (ActiveSpan -> m a) -> m (Traced a)
Tracer.traced Tracer
t SpanOpts
opt 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 = 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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall a. HasSpanFields a => Lens' a SpanContext
spanContext forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadIO m => ActiveSpan -> m Span
readActiveSpan ActiveSpan
span

        , mModifyResponse :: Response BodyReader -> IO (Response BodyReader)
mModifyResponse = \Response BodyReader
rs -> do
            forall (m :: * -> *).
MonadIO m =>
ActiveSpan -> (Span -> Span) -> m ()
modifyActiveSpan ActiveSpan
span forall a b. (a -> b) -> a -> b
$
                forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over forall a. HasSpanFields a => Lens' a Tags
spanTags (Tag -> Tags -> Tags
setTag (Status -> Tag
HttpStatusCode (forall body. Response body -> Status
responseStatus Response BodyReader
rs)))
            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 forall a. Semigroup a => a -> a -> a
<> forall c r (p :: [*]).
(HasCarrier c p, HasPropagation r p) =>
r -> SpanContext -> c
Propagation.inject Propagation p
p SpanContext
ctx
        }