-- | Module for using a WAI Middleware as an X-Ray client
module Network.AWS.XRayClient.WAI
  ( XRayClientConfig(..)
  , xrayClientConfig
  , xrayTraceMiddleware
  , NoAddressInfoException(..)
  , xrayWaiVaultKey
  , vaultDataFromRequest
  , XRayVaultData(..)
  , traceXRaySubsegment
  , traceXRaySubsegment'
  , atomicallyAddVaultDataSubsegment
  , makeSubsegmentIndependent
  , module X
  ) where

import Prelude

import Control.Concurrent (forkIO)
import Control.Lens hiding ((|>))
import Control.Monad (void, when)
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.IO.Unlift (MonadUnliftIO)
import qualified Data.ByteString.Char8 as BS8
import Data.Foldable (toList)
import Data.IORef
import Data.Sequence (Seq, (|>))
import qualified Data.Sequence as Seq
import Data.Text (Text)
import qualified Data.Text as T
import Data.Time.Clock.POSIX
import qualified Data.Vault.Lazy as V
import Network.AWS.XRayClient.Segment as X
import Network.AWS.XRayClient.SendSegments as X
import Network.AWS.XRayClient.TraceId as X
import Network.HTTP.Types.Status (statusCode)
import Network.Wai
import System.IO.Unsafe (unsafePerformIO)
import System.Random
import System.Random.XRayCustom
import UnliftIO.Exception (finally)

-- | Configuration type for the XRay client middleware.
data XRayClientConfig = XRayClientConfig
  { XRayClientConfig -> Text
xrayClientConfigDaemonHost :: !Text
  -- ^ The host that the daemon is listening on.
  , XRayClientConfig -> Int
xrayClientConfigDaemonPort :: !Int
  -- ^ The port that the daemon is listening on.
  , XRayClientConfig -> Text
xrayClientConfigApplicationName :: !Text
  -- ^ The value of the "name" field that will be sent to X-Ray.
  , XRayClientConfig
-> Maybe (Request -> Response -> POSIXTime -> POSIXTime -> IO Bool)
xrayClientConfigSampler
      :: !(Maybe (Request -> Response -> POSIXTime -> POSIXTime -> IO Bool))
  -- ^ A sampling function to filter traces.
  }

-- | Constructor for 'XRayClientConfig' with required arguments.
xrayClientConfig :: Text -> XRayClientConfig
xrayClientConfig :: Text -> XRayClientConfig
xrayClientConfig Text
appName = XRayClientConfig :: Text
-> Int
-> Text
-> Maybe (Request -> Response -> POSIXTime -> POSIXTime -> IO Bool)
-> XRayClientConfig
XRayClientConfig
  { xrayClientConfigDaemonHost :: Text
xrayClientConfigDaemonHost = Text
"127.0.0.1"
  , xrayClientConfigDaemonPort :: Int
xrayClientConfigDaemonPort = Int
2000
  , xrayClientConfigApplicationName :: Text
xrayClientConfigApplicationName = Text
appName
  , xrayClientConfigSampler :: Maybe (Request -> Response -> POSIXTime -> POSIXTime -> IO Bool)
xrayClientConfigSampler = Maybe (Request -> Response -> POSIXTime -> POSIXTime -> IO Bool)
forall a. Maybe a
Nothing
  }

-- | Traces the execution time of a request and sends the the local X-Ray
-- daemon.
xrayTraceMiddleware :: XRayClientConfig -> Middleware
xrayTraceMiddleware :: XRayClientConfig -> Middleware
xrayTraceMiddleware clientConfig :: XRayClientConfig
clientConfig@XRayClientConfig {Int
Maybe (Request -> Response -> POSIXTime -> POSIXTime -> IO Bool)
Text
xrayClientConfigSampler :: Maybe (Request -> Response -> POSIXTime -> POSIXTime -> IO Bool)
xrayClientConfigApplicationName :: Text
xrayClientConfigDaemonPort :: Int
xrayClientConfigDaemonHost :: Text
xrayClientConfigSampler :: XRayClientConfig
-> Maybe (Request -> Response -> POSIXTime -> POSIXTime -> IO Bool)
xrayClientConfigApplicationName :: XRayClientConfig -> Text
xrayClientConfigDaemonPort :: XRayClientConfig -> Int
xrayClientConfigDaemonHost :: XRayClientConfig -> Text
..} Application
app Request
request Response -> IO ResponseReceived
respond = do
  -- Start timer. We purposely include all of this setup in the time.
  POSIXTime
startTime <- IO POSIXTime
getPOSIXTime

  -- Create an IORef StdGen. We share this across the life of the request so we
  -- don't have multiple simultaneous requests causing contention on the global
  -- StdGen.
  IORef StdGen
stdGenIORef <- IO StdGen
forall (m :: * -> *). MonadIO m => m StdGen
newStdGen IO StdGen -> (StdGen -> IO (IORef StdGen)) -> IO (IORef StdGen)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= StdGen -> IO (IORef StdGen)
forall a. a -> IO (IORef a)
newIORef

  -- Check for the X-Amzn-Trace-Id and try to parse it. If that isn't possible,
  -- then make a new header.
  let
    oldHeaders :: RequestHeaders
oldHeaders = Request -> RequestHeaders
requestHeaders Request
request
    mTraceIdHeaderData :: Maybe XRayTraceIdHeaderData
mTraceIdHeaderData =
      HeaderName -> RequestHeaders -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup HeaderName
amazonTraceIdHeaderName RequestHeaders
oldHeaders Maybe ByteString
-> (ByteString -> Maybe XRayTraceIdHeaderData)
-> Maybe XRayTraceIdHeaderData
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ByteString -> Maybe XRayTraceIdHeaderData
parseXRayTraceIdHeaderData
  (RequestHeaders
newHeaders, headerData :: XRayTraceIdHeaderData
headerData@XRayTraceIdHeaderData {Maybe Bool
Maybe XRaySegmentId
XRayTraceId
xrayTraceIdHeaderDataRootTraceId :: XRayTraceIdHeaderData -> XRayTraceId
xrayTraceIdHeaderDataParentId :: XRayTraceIdHeaderData -> Maybe XRaySegmentId
xrayTraceIdHeaderDataSampled :: XRayTraceIdHeaderData -> Maybe Bool
xrayTraceIdHeaderDataSampled :: Maybe Bool
xrayTraceIdHeaderDataParentId :: Maybe XRaySegmentId
xrayTraceIdHeaderDataRootTraceId :: XRayTraceId
..}) <-
    case Maybe XRayTraceIdHeaderData
mTraceIdHeaderData of
      Just XRayTraceIdHeaderData
headerData -> (RequestHeaders, XRayTraceIdHeaderData)
-> IO (RequestHeaders, XRayTraceIdHeaderData)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RequestHeaders
oldHeaders, XRayTraceIdHeaderData
headerData)
      Maybe XRayTraceIdHeaderData
Nothing -> do
        XRayTraceId
traceId <- IORef StdGen -> IO XRayTraceId
generateXRayTraceId IORef StdGen
stdGenIORef
        let
          headerData :: XRayTraceIdHeaderData
headerData = XRayTraceId -> XRayTraceIdHeaderData
xrayTraceIdHeaderData XRayTraceId
traceId
          header :: (HeaderName, ByteString)
header =
            (HeaderName
amazonTraceIdHeaderName, XRayTraceIdHeaderData -> ByteString
makeXRayTraceIdHeaderValue XRayTraceIdHeaderData
headerData)
        (RequestHeaders, XRayTraceIdHeaderData)
-> IO (RequestHeaders, XRayTraceIdHeaderData)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((HeaderName, ByteString)
header (HeaderName, ByteString) -> RequestHeaders -> RequestHeaders
forall a. a -> [a] -> [a]
: RequestHeaders
oldHeaders, XRayTraceIdHeaderData
headerData)

  -- Construct the new request including the vault data
  XRaySegmentId
segmentId <- IORef StdGen
-> (StdGen -> (XRaySegmentId, StdGen)) -> IO XRaySegmentId
forall g a. RandomGen g => IORef g -> (g -> (a, g)) -> IO a
withRandomGenIORef IORef StdGen
stdGenIORef StdGen -> (XRaySegmentId, StdGen)
generateXRaySegmentId
  IORef (Seq XRaySegment)
subsegmentsIORef <- Seq XRaySegment -> IO (IORef (Seq XRaySegment))
forall a. a -> IO (IORef a)
newIORef Seq XRaySegment
forall a. Seq a
Seq.empty
  let
    vaultData :: XRayVaultData
vaultData = XRayTraceIdHeaderData
-> XRayClientConfig
-> XRaySegmentId
-> IORef StdGen
-> IORef (Seq XRaySegment)
-> XRayVaultData
XRayVaultData
      XRayTraceIdHeaderData
headerData
      XRayClientConfig
clientConfig
      XRaySegmentId
segmentId
      IORef StdGen
stdGenIORef
      IORef (Seq XRaySegment)
subsegmentsIORef
    request' :: Request
request' = Request
request
      { requestHeaders :: RequestHeaders
requestHeaders = RequestHeaders
newHeaders
      , vault :: Vault
vault = Key XRayVaultData -> XRayVaultData -> Vault -> Vault
forall a. Key a -> a -> Vault -> Vault
V.insert Key XRayVaultData
xrayWaiVaultKey XRayVaultData
vaultData (Request -> Vault
vault Request
request)
      }

  -- Run application with the new Request
  Application
app Request
request' ((Response -> IO ResponseReceived) -> IO ResponseReceived)
-> (Response -> IO ResponseReceived) -> IO ResponseReceived
forall a b. (a -> b) -> a -> b
$ \Response
response -> do
    ResponseReceived
responseAccept <- Response -> IO ResponseReceived
respond Response
response

    -- Create segment
    POSIXTime
endTime <- IO POSIXTime
getPOSIXTime
    Seq XRaySegment
subsegments <- IORef (Seq XRaySegment) -> IO (Seq XRaySegment)
forall a. IORef a -> IO a
readIORef IORef (Seq XRaySegment)
subsegmentsIORef
    let
      segment :: XRaySegment
segment =
        Text
-> XRaySegmentId
-> XRayTraceId
-> POSIXTime
-> Maybe POSIXTime
-> XRaySegment
xraySegment
            Text
xrayClientConfigApplicationName
            XRaySegmentId
segmentId
            XRayTraceId
xrayTraceIdHeaderDataRootTraceId
            POSIXTime
startTime
            (POSIXTime -> Maybe POSIXTime
forall a. a -> Maybe a
Just POSIXTime
endTime)
          XRaySegment -> (XRaySegment -> XRaySegment) -> XRaySegment
forall a b. a -> (a -> b) -> b
& (Maybe XRaySegmentId -> Identity (Maybe XRaySegmentId))
-> XRaySegment -> Identity XRaySegment
Lens' XRaySegment (Maybe XRaySegmentId)
xraySegmentParentId
          ((Maybe XRaySegmentId -> Identity (Maybe XRaySegmentId))
 -> XRaySegment -> Identity XRaySegment)
-> Maybe XRaySegmentId -> XRaySegment -> XRaySegment
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Maybe XRaySegmentId
xrayTraceIdHeaderDataParentId
          XRaySegment -> (XRaySegment -> XRaySegment) -> XRaySegment
forall a b. a -> (a -> b) -> b
& Request -> XRaySegment -> XRaySegment
addRequestToSegment Request
request'
          XRaySegment -> (XRaySegment -> XRaySegment) -> XRaySegment
forall a b. a -> (a -> b) -> b
& Response -> XRaySegment -> XRaySegment
addResponseToSegment Response
response
      independentSubsegments :: [XRaySegment]
independentSubsegments =
        XRayVaultData -> XRaySegment -> XRaySegment
makeSubsegmentIndependent XRayVaultData
vaultData (XRaySegment -> XRaySegment) -> [XRaySegment] -> [XRaySegment]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Seq XRaySegment -> [XRaySegment]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq XRaySegment
subsegments

    -- Send the segment if it passes the filters
    Bool
shouldSend <- IO Bool
-> ((Request -> Response -> POSIXTime -> POSIXTime -> IO Bool)
    -> IO Bool)
-> Maybe (Request -> Response -> POSIXTime -> POSIXTime -> IO Bool)
-> IO Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
      (Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True)
      (\Request -> Response -> POSIXTime -> POSIXTime -> IO Bool
f -> Request -> Response -> POSIXTime -> POSIXTime -> IO Bool
f Request
request Response
response POSIXTime
startTime POSIXTime
endTime)
      Maybe (Request -> Response -> POSIXTime -> POSIXTime -> IO Bool)
xrayClientConfigSampler
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
shouldSend (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ Text -> Int -> [XRaySegment] -> IO ()
sendSegmentsToDaemon
      Text
xrayClientConfigDaemonHost
      Int
xrayClientConfigDaemonPort
      (XRaySegment
segment XRaySegment -> [XRaySegment] -> [XRaySegment]
forall a. a -> [a] -> [a]
: [XRaySegment]
independentSubsegments)

    -- Be a good middleware and return the original ResponseAccept
    ResponseReceived -> IO ResponseReceived
forall (f :: * -> *) a. Applicative f => a -> f a
pure ResponseReceived
responseAccept

-- | Adds the info from a WAI 'Request' to an 'XRaySegment'.
addRequestToSegment :: Request -> XRaySegment -> XRaySegment
addRequestToSegment :: Request -> XRaySegment -> XRaySegment
addRequestToSegment Request
request =
  (Maybe XRaySegmentHttp -> Identity (Maybe XRaySegmentHttp))
-> XRaySegment -> Identity XRaySegment
Lens' XRaySegment (Maybe XRaySegmentHttp)
xraySegmentHttp
    ((Maybe XRaySegmentHttp -> Identity (Maybe XRaySegmentHttp))
 -> XRaySegment -> Identity XRaySegment)
-> ((XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest)
    -> Maybe XRaySegmentHttp -> Identity (Maybe XRaySegmentHttp))
-> (XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest)
-> XRaySegment
-> Identity XRaySegment
forall b c a. (b -> c) -> (a -> b) -> a -> c
. XRaySegmentHttp -> Iso' (Maybe XRaySegmentHttp) XRaySegmentHttp
forall a. Eq a => a -> Iso' (Maybe a) a
non XRaySegmentHttp
xraySegmentHttpDef
    ((XRaySegmentHttp -> Identity XRaySegmentHttp)
 -> Maybe XRaySegmentHttp -> Identity (Maybe XRaySegmentHttp))
-> ((XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest)
    -> XRaySegmentHttp -> Identity XRaySegmentHttp)
-> (XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest)
-> Maybe XRaySegmentHttp
-> Identity (Maybe XRaySegmentHttp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe XRaySegmentHttpRequest
 -> Identity (Maybe XRaySegmentHttpRequest))
-> XRaySegmentHttp -> Identity XRaySegmentHttp
Lens' XRaySegmentHttp (Maybe XRaySegmentHttpRequest)
xraySegmentHttpRequest
    ((Maybe XRaySegmentHttpRequest
  -> Identity (Maybe XRaySegmentHttpRequest))
 -> XRaySegmentHttp -> Identity XRaySegmentHttp)
-> ((XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest)
    -> Maybe XRaySegmentHttpRequest
    -> Identity (Maybe XRaySegmentHttpRequest))
-> (XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest)
-> XRaySegmentHttp
-> Identity XRaySegmentHttp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. XRaySegmentHttpRequest
-> Iso' (Maybe XRaySegmentHttpRequest) XRaySegmentHttpRequest
forall a. Eq a => a -> Iso' (Maybe a) a
non XRaySegmentHttpRequest
xraySegmentHttpRequestDef
    ((XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest)
 -> XRaySegment -> Identity XRaySegment)
-> (XRaySegmentHttpRequest -> XRaySegmentHttpRequest)
-> XRaySegment
-> XRaySegment
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (((Maybe Text -> Identity (Maybe Text))
-> XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest
Lens' XRaySegmentHttpRequest (Maybe Text)
xraySegmentHttpRequestMethod
        ((Maybe Text -> Identity (Maybe Text))
 -> XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest)
-> Text -> XRaySegmentHttpRequest -> XRaySegmentHttpRequest
forall s t a b. ASetter s t a (Maybe b) -> b -> s -> t
?~ String -> Text
T.pack (ByteString -> String
BS8.unpack (Request -> ByteString
requestMethod Request
request))
        )
       (XRaySegmentHttpRequest -> XRaySegmentHttpRequest)
-> (XRaySegmentHttpRequest -> XRaySegmentHttpRequest)
-> XRaySegmentHttpRequest
-> XRaySegmentHttpRequest
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Maybe Text -> Identity (Maybe Text))
-> XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest
Lens' XRaySegmentHttpRequest (Maybe Text)
xraySegmentHttpRequestUrl ((Maybe Text -> Identity (Maybe Text))
 -> XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest)
-> Text -> XRaySegmentHttpRequest -> XRaySegmentHttpRequest
forall s t a b. ASetter s t a (Maybe b) -> b -> s -> t
?~ Text
"/" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> [Text] -> Text
T.intercalate
           Text
"/"
           (Request -> [Text]
pathInfo Request
request)
         )
       (XRaySegmentHttpRequest -> XRaySegmentHttpRequest)
-> (XRaySegmentHttpRequest -> XRaySegmentHttpRequest)
-> XRaySegmentHttpRequest
-> XRaySegmentHttpRequest
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Maybe Text -> Identity (Maybe Text))
-> XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest
Lens' XRaySegmentHttpRequest (Maybe Text)
xraySegmentHttpRequestUserAgent
         ((Maybe Text -> Identity (Maybe Text))
 -> XRaySegmentHttpRequest -> Identity XRaySegmentHttpRequest)
-> Maybe Text -> XRaySegmentHttpRequest -> XRaySegmentHttpRequest
forall s t a b. ASetter s t a b -> b -> s -> t
.~ (String -> Text
T.pack (String -> Text) -> (ByteString -> String) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String
BS8.unpack (ByteString -> Text) -> Maybe ByteString -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Request -> Maybe ByteString
requestHeaderUserAgent Request
request)
         )
       )

-- | Adds the info from a WAI 'Response' to an 'XRaySegment'.
addResponseToSegment :: Response -> XRaySegment -> XRaySegment
addResponseToSegment :: Response -> XRaySegment -> XRaySegment
addResponseToSegment Response
response =
  (Maybe XRaySegmentHttp -> Identity (Maybe XRaySegmentHttp))
-> XRaySegment -> Identity XRaySegment
Lens' XRaySegment (Maybe XRaySegmentHttp)
xraySegmentHttp
    ((Maybe XRaySegmentHttp -> Identity (Maybe XRaySegmentHttp))
 -> XRaySegment -> Identity XRaySegment)
-> ((XRaySegmentHttpResponse -> Identity XRaySegmentHttpResponse)
    -> Maybe XRaySegmentHttp -> Identity (Maybe XRaySegmentHttp))
-> (XRaySegmentHttpResponse -> Identity XRaySegmentHttpResponse)
-> XRaySegment
-> Identity XRaySegment
forall b c a. (b -> c) -> (a -> b) -> a -> c
. XRaySegmentHttp -> Iso' (Maybe XRaySegmentHttp) XRaySegmentHttp
forall a. Eq a => a -> Iso' (Maybe a) a
non XRaySegmentHttp
xraySegmentHttpDef
    ((XRaySegmentHttp -> Identity XRaySegmentHttp)
 -> Maybe XRaySegmentHttp -> Identity (Maybe XRaySegmentHttp))
-> ((XRaySegmentHttpResponse -> Identity XRaySegmentHttpResponse)
    -> XRaySegmentHttp -> Identity XRaySegmentHttp)
-> (XRaySegmentHttpResponse -> Identity XRaySegmentHttpResponse)
-> Maybe XRaySegmentHttp
-> Identity (Maybe XRaySegmentHttp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe XRaySegmentHttpResponse
 -> Identity (Maybe XRaySegmentHttpResponse))
-> XRaySegmentHttp -> Identity XRaySegmentHttp
Lens' XRaySegmentHttp (Maybe XRaySegmentHttpResponse)
xraySegmentHttpResponse
    ((Maybe XRaySegmentHttpResponse
  -> Identity (Maybe XRaySegmentHttpResponse))
 -> XRaySegmentHttp -> Identity XRaySegmentHttp)
-> ((XRaySegmentHttpResponse -> Identity XRaySegmentHttpResponse)
    -> Maybe XRaySegmentHttpResponse
    -> Identity (Maybe XRaySegmentHttpResponse))
-> (XRaySegmentHttpResponse -> Identity XRaySegmentHttpResponse)
-> XRaySegmentHttp
-> Identity XRaySegmentHttp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. XRaySegmentHttpResponse
-> Iso' (Maybe XRaySegmentHttpResponse) XRaySegmentHttpResponse
forall a. Eq a => a -> Iso' (Maybe a) a
non XRaySegmentHttpResponse
xraySegmentHttpResponseDef
    ((XRaySegmentHttpResponse -> Identity XRaySegmentHttpResponse)
 -> XRaySegment -> Identity XRaySegment)
-> (XRaySegmentHttpResponse -> XRaySegmentHttpResponse)
-> XRaySegment
-> XRaySegment
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ ((Maybe Int -> Identity (Maybe Int))
-> XRaySegmentHttpResponse -> Identity XRaySegmentHttpResponse
Lens' XRaySegmentHttpResponse (Maybe Int)
xraySegmentHttpResponseStatus ((Maybe Int -> Identity (Maybe Int))
 -> XRaySegmentHttpResponse -> Identity XRaySegmentHttpResponse)
-> Int -> XRaySegmentHttpResponse -> XRaySegmentHttpResponse
forall s t a b. ASetter s t a (Maybe b) -> b -> s -> t
?~ Status -> Int
statusCode (Response -> Status
responseStatus Response
response))

-- | We use the WAI 'V.Vault' to store data needed during traces.
data XRayVaultData = XRayVaultData
  { XRayVaultData -> XRayTraceIdHeaderData
xrayVaultDataTraceIdHeaderData :: !XRayTraceIdHeaderData
  -- ^ Data about the current trace.
  , XRayVaultData -> XRayClientConfig
xrayVaultDataClientConfig :: !XRayClientConfig
  -- ^ Client configuration passed into the middleware.
  , XRayVaultData -> XRaySegmentId
xrayVaultDataRootSegmentId :: !XRaySegmentId
  -- ^ Segment ID of the root segment for this request.
  , XRayVaultData -> IORef StdGen
xrayVaultDataStdGen :: !(IORef StdGen)
  -- ^ 'StdGen' for generating segment IDs and trace IDs.
  , XRayVaultData -> IORef (Seq XRaySegment)
xrayVaultDataSubsegments :: !(IORef (Seq XRaySegment))
  -- ^ Current list of subsegments.
  }

-- | This is a 'V.Key' for the @vault@ inside each WAI 'Request'. It is used to
-- get to the 'XRayVaultData' for the current request.
xrayWaiVaultKey :: V.Key XRayVaultData
xrayWaiVaultKey :: Key XRayVaultData
xrayWaiVaultKey = IO (Key XRayVaultData) -> Key XRayVaultData
forall a. IO a -> a
unsafePerformIO IO (Key XRayVaultData)
forall a. IO (Key a)
V.newKey
{-# NOINLINE xrayWaiVaultKey #-}

-- | Try to get 'XRayVaultData' from the WAI 'Request' vault.
vaultDataFromRequest :: Request -> Maybe XRayVaultData
vaultDataFromRequest :: Request -> Maybe XRayVaultData
vaultDataFromRequest = Key XRayVaultData -> Vault -> Maybe XRayVaultData
forall a. Key a -> Vault -> Maybe a
V.lookup Key XRayVaultData
xrayWaiVaultKey (Vault -> Maybe XRayVaultData)
-> (Request -> Vault) -> Request -> Maybe XRayVaultData
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Request -> Vault
vault

-- | Time a 'MonadIO' action and add it to the list of subsegments.
traceXRaySubsegment
  :: MonadUnliftIO m
  => Request
  -> Text
  -> (XRaySegment -> XRaySegment)
  -> m a
  -> m a
traceXRaySubsegment :: Request -> Text -> (XRaySegment -> XRaySegment) -> m a -> m a
traceXRaySubsegment Request
request Text
subsegmentName XRaySegment -> XRaySegment
modifySubsegment m a
action =
  case Key XRayVaultData -> Vault -> Maybe XRayVaultData
forall a. Key a -> Vault -> Maybe a
V.lookup Key XRayVaultData
xrayWaiVaultKey (Request -> Vault
vault Request
request) of
    Maybe XRayVaultData
Nothing -> m a
action
    Just XRayVaultData
v -> XRayVaultData -> Text -> (XRaySegment -> XRaySegment) -> m a -> m a
forall (m :: * -> *) a.
MonadUnliftIO m =>
XRayVaultData -> Text -> (XRaySegment -> XRaySegment) -> m a -> m a
traceXRaySubsegment' XRayVaultData
v Text
subsegmentName XRaySegment -> XRaySegment
modifySubsegment m a
action

traceXRaySubsegment'
  :: MonadUnliftIO m
  => XRayVaultData
  -> Text
  -> (XRaySegment -> XRaySegment)
  -> m a
  -> m a
traceXRaySubsegment' :: XRayVaultData -> Text -> (XRaySegment -> XRaySegment) -> m a -> m a
traceXRaySubsegment' XRayVaultData
vaultData Text
subsegmentName XRaySegment -> XRaySegment
modifySubsegment m a
action = do
  -- Run action with timing
  POSIXTime
startTime <- IO POSIXTime -> m POSIXTime
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO POSIXTime
getPOSIXTime

  -- Catch any exceptions and rethrow them once we've sent the segment
  m a -> m () -> m a
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
finally m a
action (m () -> m a) -> m () -> m a
forall a b. (a -> b) -> a -> b
$ IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    POSIXTime
endTime <- IO POSIXTime
getPOSIXTime

    XRaySegmentId
segmentId <- IORef StdGen
-> (StdGen -> (XRaySegmentId, StdGen)) -> IO XRaySegmentId
forall g a. RandomGen g => IORef g -> (g -> (a, g)) -> IO a
withRandomGenIORef
      (XRayVaultData -> IORef StdGen
xrayVaultDataStdGen XRayVaultData
vaultData)
      StdGen -> (XRaySegmentId, StdGen)
generateXRaySegmentId

    let
      subsegment :: XRaySegment
subsegment =
        Text
-> XRaySegmentId -> POSIXTime -> Maybe POSIXTime -> XRaySegment
xraySubsegment Text
subsegmentName XRaySegmentId
segmentId POSIXTime
startTime (POSIXTime -> Maybe POSIXTime
forall a. a -> Maybe a
Just POSIXTime
endTime)
      subsegment' :: XRaySegment
subsegment' = XRaySegment -> XRaySegment
modifySubsegment XRaySegment
subsegment
    XRayVaultData -> XRaySegment -> IO ()
atomicallyAddVaultDataSubsegment XRayVaultData
vaultData XRaySegment
subsegment'

-- | Add subsegment to XRay vault data 'IORef'.
atomicallyAddVaultDataSubsegment :: XRayVaultData -> XRaySegment -> IO ()
atomicallyAddVaultDataSubsegment :: XRayVaultData -> XRaySegment -> IO ()
atomicallyAddVaultDataSubsegment XRayVaultData
vaultData XRaySegment
subsegment =
  IORef (Seq XRaySegment)
-> (Seq XRaySegment -> (Seq XRaySegment, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (XRayVaultData -> IORef (Seq XRaySegment)
xrayVaultDataSubsegments XRayVaultData
vaultData)
    ((Seq XRaySegment -> (Seq XRaySegment, ())) -> IO ())
-> (Seq XRaySegment -> (Seq XRaySegment, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Seq XRaySegment
subsegments -> (Seq XRaySegment
subsegments Seq XRaySegment -> XRaySegment -> Seq XRaySegment
forall a. Seq a -> a -> Seq a
|> XRaySegment
subsegment, ())

-- | Uses the trace ID and segment ID of the root segment from the vault to
-- make a subsegment independent. This is useful so nested components that
-- create subsegments don't need all of this information threaded down to them.
-- We can just decorate all of the subsegments with it before sending them off.
makeSubsegmentIndependent :: XRayVaultData -> XRaySegment -> XRaySegment
makeSubsegmentIndependent :: XRayVaultData -> XRaySegment -> XRaySegment
makeSubsegmentIndependent XRayVaultData {XRaySegmentId
XRayTraceIdHeaderData
IORef StdGen
IORef (Seq XRaySegment)
XRayClientConfig
xrayVaultDataSubsegments :: IORef (Seq XRaySegment)
xrayVaultDataStdGen :: IORef StdGen
xrayVaultDataRootSegmentId :: XRaySegmentId
xrayVaultDataClientConfig :: XRayClientConfig
xrayVaultDataTraceIdHeaderData :: XRayTraceIdHeaderData
xrayVaultDataSubsegments :: XRayVaultData -> IORef (Seq XRaySegment)
xrayVaultDataStdGen :: XRayVaultData -> IORef StdGen
xrayVaultDataRootSegmentId :: XRayVaultData -> XRaySegmentId
xrayVaultDataClientConfig :: XRayVaultData -> XRayClientConfig
xrayVaultDataTraceIdHeaderData :: XRayVaultData -> XRayTraceIdHeaderData
..} XRaySegment
subsegment =
  XRaySegment
subsegment
    XRaySegment -> (XRaySegment -> XRaySegment) -> XRaySegment
forall a b. a -> (a -> b) -> b
& (Maybe XRayTraceId -> Identity (Maybe XRayTraceId))
-> XRaySegment -> Identity XRaySegment
Lens' XRaySegment (Maybe XRayTraceId)
xraySegmentTraceId
    ((Maybe XRayTraceId -> Identity (Maybe XRayTraceId))
 -> XRaySegment -> Identity XRaySegment)
-> XRayTraceId -> XRaySegment -> XRaySegment
forall s t a b. ASetter s t a (Maybe b) -> b -> s -> t
?~ XRayTraceIdHeaderData -> XRayTraceId
xrayTraceIdHeaderDataRootTraceId XRayTraceIdHeaderData
xrayVaultDataTraceIdHeaderData
    XRaySegment -> (XRaySegment -> XRaySegment) -> XRaySegment
forall a b. a -> (a -> b) -> b
& (Maybe XRaySegmentId -> Identity (Maybe XRaySegmentId))
-> XRaySegment -> Identity XRaySegment
Lens' XRaySegment (Maybe XRaySegmentId)
xraySegmentParentId
    ((Maybe XRaySegmentId -> Identity (Maybe XRaySegmentId))
 -> XRaySegment -> Identity XRaySegment)
-> XRaySegmentId -> XRaySegment -> XRaySegment
forall s t a b. ASetter s t a (Maybe b) -> b -> s -> t
?~ XRaySegmentId
xrayVaultDataRootSegmentId
    XRaySegment -> (XRaySegment -> XRaySegment) -> XRaySegment
forall a b. a -> (a -> b) -> b
& (Maybe Text -> Identity (Maybe Text))
-> XRaySegment -> Identity XRaySegment
Lens' XRaySegment (Maybe Text)
xraySegmentType
    ((Maybe Text -> Identity (Maybe Text))
 -> XRaySegment -> Identity XRaySegment)
-> Text -> XRaySegment -> XRaySegment
forall s t a b. ASetter s t a (Maybe b) -> b -> s -> t
?~ Text
"subsegment"