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

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

import           Control.Lens            (over, set, view)
import           Data.Maybe
import           Data.Semigroup
import           Data.Text               (Text)
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)


-- | A 'TracedApplication' is a WAI 'Application' with an 'ActiveSpan`.
--
-- Expanded:
--
-- @
-- type TracedApplication =
--     ActiveSpan -> Request -> (Response -> IO ResponseReceived) -> IO ResponseReceived
-- @
type TracedApplication = ActiveSpan -> Application

-- | The operation name is, basically, the name of the span.
--
-- This is typically determined from the request in some way, see
-- 'defaultOperationName'.
--
-- @since 0.2.0
type OperationName = Request -> Text

-- | Middleware to enable tracing for a WAI application.
--
-- This uses the 'defaultOperationName'.
opentracing
    :: HasCarrier Headers p
    => Tracer
    -> Propagation        p
    -> TracedApplication
    -> Application
opentracing :: forall (p :: [*]).
HasCarrier Headers p =>
Tracer -> Propagation p -> TracedApplication -> Application
opentracing Tracer
t Propagation p
p TracedApplication
app Request
req Response -> IO ResponseReceived
respond =
    forall (p :: [*]).
HasCarrier Headers p =>
Tracer
-> Propagation p
-> OperationName
-> TracedApplication
-> Application
withOperationName Tracer
t Propagation p
p OperationName
defaultOperationName TracedApplication
app Request
req Response -> IO ResponseReceived
respond

-- | Customise the tracing middleware with an 'OperationName'.
--
-- It is intended to import this module qualified for legibility
-- (@OpenTracing.withOperationName@).
--
-- @since 0.2.0
withOperationName
    :: HasCarrier Headers p
    => Tracer
    -> Propagation        p
    -> OperationName
    -> TracedApplication
    -> Application
withOperationName :: forall (p :: [*]).
HasCarrier Headers p =>
Tracer
-> Propagation p
-> OperationName
-> TracedApplication
-> Application
withOperationName Tracer
t Propagation p
p OperationName
opname TracedApplication
app Request
req Response -> IO ResponseReceived
respond = do
    let ctx :: Maybe SpanContext
ctx = forall c r (p :: [*]).
(HasCarrier c p, HasPropagation r p) =>
r -> c -> Maybe SpanContext
Propagation.extract Propagation p
p (Request -> Headers
requestHeaders Request
req)
    let opt :: SpanOpts
opt = let name :: Text
name = OperationName
opname Request
req
                  refs :: SpanRefs
refs = (\[Reference]
x -> forall s t a b. ASetter s t a b -> b -> s -> t
set Lens' SpanRefs [Reference]
refPropagated [Reference]
x forall a. Monoid a => a
mempty)
                       forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Maybe a -> [a]
maybeToList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SpanContext -> Reference
ChildOf forall a b. (a -> b) -> a -> b
$ Maybe SpanContext
ctx
               in forall s t a b. ASetter s t a b -> b -> s -> t
set Lens' SpanOpts (Maybe Sampled)
spanOptSampled (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' SpanContext Sampled
ctxSampled forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe SpanContext
ctx)
                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
requestMethod Request
req)
                      , Text -> Tag
HttpUrl     (Method -> Text
decodeUtf8 Method
url)
                      , Text -> Tag
PeerAddress (String -> Text
Text.pack (forall a. Show a => a -> String
show (Request -> SockAddr
remoteHost Request
req))) -- not so great
                      , SpanKinds -> Tag
SpanKind    SpanKinds
RPCServer
                      ]
                forall a b. (a -> b) -> a -> b
$ Text -> SpanRefs -> SpanOpts
spanOpts Text
name SpanRefs
refs

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

-- | The default 'OperationName' is the @pathInfo@ of the request.
--
-- @since 0.2.0
defaultOperationName :: OperationName
defaultOperationName :: OperationName
defaultOperationName Request
req = Char -> Text -> Text
Text.cons Char
'/' (Text -> [Text] -> Text
Text.intercalate Text
"/" (Request -> [Text]
pathInfo Request
req))