diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,7 @@ +## [*Unreleased*](https://github.com/freckle/aws-xray-client/tree/aws-xray-client-wai-v0.1.0.0...main) + +None + +## [v0.1.0.0](https://github.com/freckle/aws-xray-client/tree/aws-xray-client-wai-v0.1.0.0) + +First tagged release. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,7 @@ +# aws-xray-client-wai + +A Haskell client for integrating AWS X-Ray with [WAI](https://hackage.haskell.org/package/wai). + +To configure this, you will need to configure a `import Network.AWS.XRayClient.WAI.XRayClientConfig)`. You can use `xrayClientConfig :: Text -> XRayClientConfig` to create a default configuration, given just an app name. + +`xrayTraceMiddleware` then accepts `XRayClientConfig` to create [WAI middleware](https://www.yesodweb.com/book/web-application-interface#web-application-interface_middleware). diff --git a/aws-xray-client-wai.cabal b/aws-xray-client-wai.cabal new file mode 100644 --- /dev/null +++ b/aws-xray-client-wai.cabal @@ -0,0 +1,50 @@ +cabal-version: 1.12 +name: aws-xray-client-wai +version: 0.1.0.0 +license: MIT +copyright: 2021 Renaissance Learning Inc +maintainer: engineering@freckle.com +author: Freckle R&D +homepage: https://github.com/freckle/aws-xray-client#readme +bug-reports: https://github.com/freckle/aws-xray-client/issues +synopsis: A client for AWS X-Ray integration with WAI. +description: + Works with `aws-xray-client` to enable X-Ray tracing with WAI. + +build-type: Simple +extra-source-files: + README.md + CHANGELOG.md + +source-repository head + type: git + location: https://github.com/freckle/aws-xray-client + +library + exposed-modules: Network.AWS.XRayClient.WAI + hs-source-dirs: library + other-modules: Paths_aws_xray_client_wai + default-language: Haskell2010 + default-extensions: + BangPatterns DataKinds DeriveAnyClass DeriveFoldable DeriveFunctor + DeriveGeneric DeriveLift DeriveTraversable DerivingStrategies + FlexibleContexts FlexibleInstances GADTs GeneralizedNewtypeDeriving + LambdaCase MultiParamTypeClasses NoImplicitPrelude + NoMonomorphismRestriction OverloadedStrings RankNTypes + RecordWildCards ScopedTypeVariables StandaloneDeriving + TypeApplications TypeFamilies + + build-depends: + aws-xray-client >=0.1.0.0, + base >=4.12.0.0 && <5, + bytestring >=0.10.8.2, + containers >=0.6.0.1, + http-types >=0.12.3, + lens >=4.17.1, + random >=1.1, + text >=1.2.3.1, + time >=1.8.0.2, + unliftio >=0.2.11, + unliftio-core >=0.1.2.0, + vault >=0.3.1.2, + wai >=3.2.2.1 diff --git a/library/Network/AWS/XRayClient/WAI.hs b/library/Network/AWS/XRayClient/WAI.hs new file mode 100644 --- /dev/null +++ b/library/Network/AWS/XRayClient/WAI.hs @@ -0,0 +1,261 @@ +{-# LANGUAGE NamedFieldPuns #-} + +-- | 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.Monoid ((<>)) +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 + { xrayClientConfigDaemonHost :: !Text + -- ^ The host that the daemon is listening on. + , xrayClientConfigDaemonPort :: !Int + -- ^ The port that the daemon is listening on. + , xrayClientConfigApplicationName :: !Text + -- ^ The value of the "name" field that will be sent to X-Ray. + , xrayClientConfigSampler :: !(Maybe (Request -> Response -> POSIXTime -> POSIXTime -> IO Bool)) + -- ^ A sampling function to filter traces. + } + +-- | Constructor for 'XRayClientConfig' with required arguments. +xrayClientConfig :: Text -> XRayClientConfig +xrayClientConfig appName = XRayClientConfig + { xrayClientConfigDaemonHost = "127.0.0.1" + , xrayClientConfigDaemonPort = 2000 + , xrayClientConfigApplicationName = appName + , xrayClientConfigSampler = Nothing + } + +-- | Traces the execution time of a request and sends the the local X-Ray +-- daemon. +xrayTraceMiddleware :: XRayClientConfig -> Middleware +xrayTraceMiddleware clientConfig@XRayClientConfig {..} app request respond = do + -- Start timer. We purposely include all of this setup in the time. + startTime <- 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. + stdGenIORef <- newStdGen >>= 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 request + mTraceIdHeaderData = + lookup amazonTraceIdHeaderName oldHeaders >>= parseXRayTraceIdHeaderData + (newHeaders, headerData@XRayTraceIdHeaderData {..}) <- + case mTraceIdHeaderData of + Just headerData -> pure (oldHeaders, headerData) + Nothing -> do + traceId <- generateXRayTraceId stdGenIORef + let + headerData = xrayTraceIdHeaderData traceId + header = + (amazonTraceIdHeaderName, makeXRayTraceIdHeaderValue headerData) + pure (header : oldHeaders, headerData) + + -- Construct the new request including the vault data + segmentId <- withRandomGenIORef stdGenIORef generateXRaySegmentId + subsegmentsIORef <- newIORef Seq.empty + let + vaultData = XRayVaultData + headerData + clientConfig + segmentId + stdGenIORef + subsegmentsIORef + request' = request + { requestHeaders = newHeaders + , vault = V.insert xrayWaiVaultKey vaultData (vault request) + } + + -- Run application with the new Request + app request' $ \response -> do + responseAccept <- respond response + + -- Create segment + endTime <- getPOSIXTime + subsegments <- readIORef subsegmentsIORef + let + segment = + xraySegment + xrayClientConfigApplicationName + segmentId + xrayTraceIdHeaderDataRootTraceId + startTime + (Just endTime) + & xraySegmentParentId + .~ xrayTraceIdHeaderDataParentId + & addRequestToSegment request' + & addResponseToSegment response + independentSubsegments = + makeSubsegmentIndependent vaultData <$> toList subsegments + + -- Send the segment if it passes the filters + shouldSend <- maybe + (pure True) + (\f -> f request response startTime endTime) + xrayClientConfigSampler + when shouldSend $ void $ forkIO $ sendSegmentsToDaemon + xrayClientConfigDaemonHost + xrayClientConfigDaemonPort + (segment : independentSubsegments) + + -- Be a good middleware and return the original ResponseAccept + pure responseAccept + +-- | Adds the info from a WAI 'Request' to an 'XRaySegment'. +addRequestToSegment :: Request -> XRaySegment -> XRaySegment +addRequestToSegment request = + xraySegmentHttp + . non xraySegmentHttpDef + . xraySegmentHttpRequest + . non xraySegmentHttpRequestDef + %~ ((xraySegmentHttpRequestMethod + ?~ T.pack (BS8.unpack (requestMethod request)) + ) + . (xraySegmentHttpRequestUrl ?~ "/" <> T.intercalate + "/" + (pathInfo request) + ) + . (xraySegmentHttpRequestUserAgent + .~ (T.pack . BS8.unpack <$> requestHeaderUserAgent request) + ) + ) + +-- | Adds the info from a WAI 'Response' to an 'XRaySegment'. +addResponseToSegment :: Response -> XRaySegment -> XRaySegment +addResponseToSegment response = + xraySegmentHttp + . non xraySegmentHttpDef + . xraySegmentHttpResponse + . non xraySegmentHttpResponseDef + %~ (xraySegmentHttpResponseStatus ?~ statusCode (responseStatus response)) + +-- | We use the WAI 'V.Vault' to store data needed during traces. +data XRayVaultData + = XRayVaultData + { xrayVaultDataTraceIdHeaderData :: !XRayTraceIdHeaderData + -- ^ Data about the current trace. + , xrayVaultDataClientConfig :: !XRayClientConfig + -- ^ Client configuration passed into the middleware. + , xrayVaultDataRootSegmentId :: !XRaySegmentId + -- ^ Segment ID of the root segment for this request. + , xrayVaultDataStdGen :: !(IORef StdGen) + -- ^ 'StdGen' for generating segment IDs and trace IDs. + , 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 = unsafePerformIO V.newKey +{-# NOINLINE xrayWaiVaultKey #-} + +-- | Try to get 'XRayVaultData' from the WAI 'Request' vault. +vaultDataFromRequest :: Request -> Maybe XRayVaultData +vaultDataFromRequest = V.lookup xrayWaiVaultKey . 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 subsegmentName modifySubsegment action = + case V.lookup xrayWaiVaultKey (vault request) of + Nothing -> action + Just v -> traceXRaySubsegment' v subsegmentName modifySubsegment action + +traceXRaySubsegment' + :: MonadUnliftIO m + => XRayVaultData + -> Text + -> (XRaySegment -> XRaySegment) + -> m a + -> m a +traceXRaySubsegment' vaultData subsegmentName modifySubsegment action = do + -- Run action with timing + startTime <- liftIO getPOSIXTime + + -- Catch any exceptions and rethrow them once we've sent the segment + finally action + $ liftIO + $ do + endTime <- getPOSIXTime + + -- Generate a segment ID + segmentId <- withRandomGenIORef + (xrayVaultDataStdGen vaultData) + generateXRaySegmentId + + -- Create the subsegment value and append to the subsegments. + let + subsegment = + xraySubsegment subsegmentName segmentId startTime (Just endTime) + subsegment' = modifySubsegment subsegment + atomicallyAddVaultDataSubsegment vaultData subsegment' + +-- | Add subsegment to XRay vault data 'IORef'. +atomicallyAddVaultDataSubsegment :: XRayVaultData -> XRaySegment -> IO () +atomicallyAddVaultDataSubsegment vaultData subsegment = + atomicModifyIORef' (xrayVaultDataSubsegments vaultData) + $ \subsegments -> (subsegments |> 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 {..} subsegment = + subsegment + & xraySegmentTraceId + ?~ xrayTraceIdHeaderDataRootTraceId xrayVaultDataTraceIdHeaderData + & xraySegmentParentId + ?~ xrayVaultDataRootSegmentId + & xraySegmentType + ?~ "subsegment"