module Network.Wai.Middleware.Stats
  ( addThreadContextFromStatsTags
  , requestStats
  ) where

import Freckle.App.Prelude

import Blammo.Logging (Pair, withThreadContext)
import Control.Lens ((^.))
import Control.Monad.Reader (runReaderT)
import Data.Aeson ((.=))
import qualified Freckle.App.Aeson as Key
import Freckle.App.Stats (HasStatsClient (..), tagsL)
import qualified Freckle.App.Stats as Stats
import Network.HTTP.Types.Status (Status (..))
import Network.Wai (Middleware, Request, requestMethod, responseStatus)

-- | Add any tags in the ambient 'StatsClient' to the logging context
addThreadContextFromStatsTags :: HasStatsClient env => env -> Middleware
addThreadContextFromStatsTags :: forall env. HasStatsClient env => env -> Middleware
addThreadContextFromStatsTags env
env Application
app Request
req Response -> IO ResponseReceived
respond = do
  let context :: [Pair]
context = (Text -> Text -> Pair) -> (Text, Text) -> Pair
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Text -> Text -> Pair
fromTag ((Text, Text) -> Pair) -> [(Text, Text)] -> [Pair]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> env
env env -> Getting [(Text, Text)] env [(Text, Text)] -> [(Text, Text)]
forall s a. s -> Getting a s a -> a
^. (StatsClient -> Const [(Text, Text)] StatsClient)
-> env -> Const [(Text, Text)] env
forall env. HasStatsClient env => Lens' env StatsClient
Lens' env StatsClient
statsClientL ((StatsClient -> Const [(Text, Text)] StatsClient)
 -> env -> Const [(Text, Text)] env)
-> (([(Text, Text)] -> Const [(Text, Text)] [(Text, Text)])
    -> StatsClient -> Const [(Text, Text)] StatsClient)
-> Getting [(Text, Text)] env [(Text, Text)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(Text, Text)] -> Const [(Text, Text)] [(Text, Text)])
-> StatsClient -> Const [(Text, Text)] StatsClient
Lens' StatsClient [(Text, Text)]
tagsL
  [Pair] -> IO ResponseReceived -> IO ResponseReceived
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
[Pair] -> m a -> m a
withThreadContext [Pair]
context (IO ResponseReceived -> IO ResponseReceived)
-> IO ResponseReceived -> IO ResponseReceived
forall a b. (a -> b) -> a -> b
$ Application
app Request
req Response -> IO ResponseReceived
respond
 where
  fromTag :: Text -> Text -> Pair
  fromTag :: Text -> Text -> Pair
fromTag Text
k Text
v = Text -> Key
Key.fromText Text
k Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= Text
v

-- | Emit @requests@ and @response_time_ms@ metrics
requestStats
  :: HasStatsClient env => env -> (Request -> [(Text, Text)]) -> Middleware
requestStats :: forall env.
HasStatsClient env =>
env -> (Request -> [(Text, Text)]) -> Middleware
requestStats env
env Request -> [(Text, Text)]
getTags Application
app Request
req Response -> IO ResponseReceived
respond = do
  UTCTime
start <- IO UTCTime
getCurrentTime
  Application
app Request
req ((Response -> IO ResponseReceived) -> IO ResponseReceived)
-> (Response -> IO ResponseReceived) -> IO ResponseReceived
forall a b. (a -> b) -> a -> b
$ \Response
res -> do
    let tags :: [(Text, Text)]
tags =
          Request -> [(Text, Text)]
getTags Request
req
            [(Text, Text)] -> [(Text, Text)] -> [(Text, Text)]
forall a. Semigroup a => a -> a -> a
<> [ (Text
"method", ByteString -> Text
decodeUtf8 (ByteString -> Text) -> ByteString -> Text
forall a b. (a -> b) -> a -> b
$ Request -> ByteString
requestMethod Request
req)
               , (Text
"status", String -> Text
pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show (Int -> String) -> Int -> String
forall a b. (a -> b) -> a -> b
$ Status -> Int
statusCode (Status -> Int) -> Status -> Int
forall a b. (a -> b) -> a -> b
$ Response -> Status
responseStatus Response
res)
               ]

    (ReaderT env IO () -> env -> IO ())
-> env -> ReaderT env IO () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip ReaderT env IO () -> env -> IO ()
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT env
env (ReaderT env IO () -> IO ()) -> ReaderT env IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [(Text, Text)] -> ReaderT env IO () -> ReaderT env IO ()
forall env (m :: * -> *) a.
(MonadReader env m, HasStatsClient env) =>
[(Text, Text)] -> m a -> m a
Stats.tagged [(Text, Text)]
tags (ReaderT env IO () -> ReaderT env IO ())
-> ReaderT env IO () -> ReaderT env IO ()
forall a b. (a -> b) -> a -> b
$ do
      Text -> ReaderT env IO ()
forall (m :: * -> *) env.
(MonadUnliftIO m, MonadReader env m, HasStatsClient env) =>
Text -> m ()
Stats.increment Text
"requests"
      Text -> UTCTime -> ReaderT env IO ()
forall (m :: * -> *) env.
(MonadUnliftIO m, MonadReader env m, HasStatsClient env) =>
Text -> UTCTime -> m ()
Stats.histogramSinceMs Text
"response_time_ms" UTCTime
start

    Response -> IO ResponseReceived
respond Response
res