{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns        #-}
{-# LANGUAGE OverloadedStrings     #-}

module Network.Wai.Middleware.OpenTracing
    ( TracedApplication
    , opentracing
    )
where

import           Control.Lens            (over, set, view)
import           Data.Maybe
import           Data.Semigroup
import qualified Data.Text               as Text
import           Data.Text.Encoding      (decodeUtf8)
import           Network.Wai
import           OpenTracing
import qualified OpenTracing.Propagation as Propagation
import qualified OpenTracing.Tracer      as Tracer
import           Prelude                 hiding (span)


type TracedApplication = ActiveSpan -> Application

opentracing
    :: HasCarrier Headers p
    => Tracer
    -> Propagation        p
    -> TracedApplication
    -> Application
opentracing :: Tracer -> Propagation p -> TracedApplication -> Application
opentracing Tracer
t Propagation p
p TracedApplication
app Request
req Response -> IO ResponseReceived
respond = do
    let ctx :: Maybe SpanContext
ctx = Propagation p -> RequestHeaders -> Maybe SpanContext
forall c r (p :: [*]).
(HasCarrier c p, HasPropagation r p) =>
r -> c -> Maybe SpanContext
Propagation.extract Propagation p
p (Request -> RequestHeaders
requestHeaders Request
req)
    let opt :: SpanOpts
opt = let name :: Text
name = Text -> [Text] -> Text
Text.intercalate Text
"/" (Request -> [Text]
pathInfo Request
req)
                  refs :: SpanRefs
refs = (\[Reference]
x -> ASetter SpanRefs SpanRefs [Reference] [Reference]
-> [Reference] -> SpanRefs -> SpanRefs
forall s t a b. ASetter s t a b -> b -> s -> t
set ASetter SpanRefs SpanRefs [Reference] [Reference]
Lens' SpanRefs [Reference]
refPropagated [Reference]
x SpanRefs
forall a. Monoid a => a
mempty)
                       ([Reference] -> SpanRefs)
-> (Maybe SpanContext -> [Reference])
-> Maybe SpanContext
-> SpanRefs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Reference -> [Reference]
forall a. Maybe a -> [a]
maybeToList (Maybe Reference -> [Reference])
-> (Maybe SpanContext -> Maybe Reference)
-> Maybe SpanContext
-> [Reference]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SpanContext -> Reference) -> Maybe SpanContext -> Maybe Reference
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SpanContext -> Reference
ChildOf (Maybe SpanContext -> SpanRefs) -> Maybe SpanContext -> SpanRefs
forall a b. (a -> b) -> a -> b
$ Maybe SpanContext
ctx
               in 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 (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) -> Maybe SpanContext -> Maybe Sampled
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe SpanContext
ctx)
                (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
requestMethod Request
req)
                      , Text -> Tag
HttpUrl     (Method -> Text
decodeUtf8 Method
url)
                      , Text -> Tag
PeerAddress (String -> Text
Text.pack (SockAddr -> String
forall a. Show a => a -> String
show (Request -> SockAddr
remoteHost Request
req))) -- not so great
                      , SpanKinds -> Tag
SpanKind    SpanKinds
RPCServer
                      ]
                (SpanOpts -> SpanOpts) -> SpanOpts -> SpanOpts
forall a b. (a -> b) -> a -> b
$ Text -> SpanRefs -> SpanOpts
spanOpts Text
name SpanRefs
refs

    Tracer
-> SpanOpts
-> (ActiveSpan -> IO ResponseReceived)
-> IO ResponseReceived
forall t (m :: * -> *) a.
(HasTracer t, MonadMask m, MonadIO m) =>
t -> SpanOpts -> (ActiveSpan -> m a) -> m a
Tracer.traced_ Tracer
t SpanOpts
opt ((ActiveSpan -> IO ResponseReceived) -> IO ResponseReceived)
-> (ActiveSpan -> IO ResponseReceived) -> IO ResponseReceived
forall a b. (a -> b) -> a -> b
$ \ActiveSpan
span -> TracedApplication
app ActiveSpan
span Request
req ((Response -> IO ResponseReceived) -> IO ResponseReceived)
-> (Response -> IO ResponseReceived) -> IO ResponseReceived
forall a b. (a -> b) -> a -> b
$ \Response
res -> 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 -> Status
responseStatus Response
res)))
        Response -> IO ResponseReceived
respond Response
res
  where
    url :: Method
url = Method
"http" Method -> Method -> Method
forall a. Semigroup a => a -> a -> a
<> if Request -> Bool
isSecure Request
req then Method
"s" else Method
forall a. Monoid a => a
mempty Method -> Method -> Method
forall a. Semigroup a => a -> a -> a
<> Method
"://"
       Method -> Method -> Method
forall a. Semigroup a => a -> a -> a
<> Method -> Maybe Method -> Method
forall a. a -> Maybe a -> a
fromMaybe Method
"localhost" (Request -> Maybe Method
requestHeaderHost Request
req)
       Method -> Method -> Method
forall a. Semigroup a => a -> a -> a
<> Request -> Method
rawPathInfo Request
req Method -> Method -> Method
forall a. Semigroup a => a -> a -> a
<> Request -> Method
rawQueryString Request
req