Merge pull request #1 from anchor/servant-auto

Automatic metering for servant apis

ekg-servant.cabal

@@ -16,9 +16,29 @@ source-repository HEAD
 
 library
   exposed-modules:     Servant.Ekg
+  hs-source-dirs:      lib
   build-depends:       base >=4.7 && <4.9
                      , ekg-core
+                     , servant >=0.2 && <0.3
                      , http-types
+                     , text
                      , time
+                     , unordered-containers
                      , wai
-  default-language:    Haskell2010
+  default-language:    Haskell2010
+
+test-suite spec
+  type: exitcode-stdio-1.0
+  ghc-options: -Wall
+  default-language: Haskell2010
+  hs-source-dirs: test
+  main-is: test.hs
+  build-depends:      base == 4.*
+                    , aeson
+                    , ekg
+                    , ekg-core
+                    , ekg-servant
+                    , servant-server
+                    , text
+                    , wai
+                    , warp

lib/Servant/Ekg.hs

@@ -1,14 +1,31 @@
+{-# LANGUAGE DataKinds           #-}
+{-# LANGUAGE FlexibleContexts    #-}
+{-# LANGUAGE FlexibleInstances   #-}
+{-# LANGUAGE KindSignatures      #-}
+{-# LANGUAGE OverloadedStrings   #-}
+{-# LANGUAGE RecordWildCards     #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeOperators       #-}
 module Servant.Ekg where
 
+import           Control.Concurrent.MVar
 import           Control.Exception
+import qualified Data.HashMap.Strict         as H
+import           Data.Monoid
+import           Data.Proxy
+import           Data.Text                   (Text)
+import qualified Data.Text                   as T
+import qualified Data.Text.Encoding          as T
 import           Data.Time.Clock
-import           Network.HTTP.Types
+import           GHC.TypeLits
+import           Network.HTTP.Types          (Method, Status (..))
 import           Network.Wai
+import           Servant.API
+import           System.Metrics
 import qualified System.Metrics.Counter      as Counter
 import qualified System.Metrics.Distribution as Distribution
 import qualified System.Metrics.Gauge        as Gauge
 
-
 gaugeInflight :: Gauge.Gauge -> Middleware
 gaugeInflight inflight application request respond =
     bracket_ (Gauge.inc inflight)
@@ -36,4 +53,102 @@ responseTimeDistribution dist application request respond =
     stop t1 = do
         t2 <- getCurrentTime
         let dt = diffUTCTime t2 t1
-        Distribution.add dist $ fromRational $ toRational dt
+        Distribution.add dist $ fromRational $ (*1000) $ toRational dt
+
+data Meters = Meters
+    { metersInflight :: Gauge.Gauge
+    , metersC2XX     :: Counter.Counter
+    , metersC4XX     :: Counter.Counter
+    , metersC5XX     :: Counter.Counter
+    , metersCXXX     :: Counter.Counter
+    , metersTime     :: Distribution.Distribution
+    }
+
+monitorEndpoints
+    :: HasEndpoint api
+    => Proxy api
+    -> Store
+    -> MVar (H.HashMap Text Meters)
+    -> Middleware
+monitorEndpoints proxy store meters application request respond = do
+    let path = case getEndpoint proxy request of
+            Nothing -> "unknown"
+            Just (ps,method) -> T.intercalate "." $ ps <> [T.decodeUtf8 method]
+    Meters{..} <- modifyMVar meters $ \ms -> case H.lookup path ms of
+        Nothing -> do
+            let prefix = "servant.path." <> path <> "."
+            metersInflight <- createGauge (prefix <> "in_flight") store
+            metersC2XX <- createCounter (prefix <> "responses.2XX") store
+            metersC4XX <- createCounter (prefix <> "responses.4XX") store
+            metersC5XX <- createCounter (prefix <> "responses.5XX") store
+            metersCXXX <- createCounter (prefix <> "responses.XXX") store
+            metersTime <- createDistribution (prefix <> "time_ms") store
+            let m = Meters{..}
+            return (H.insert path m ms, m)
+        Just m -> return (ms,m)
+    let application' =
+            responseTimeDistribution metersTime .
+            countResponseCodes (metersC2XX, metersC4XX, metersC5XX, metersCXXX) .
+            gaugeInflight metersInflight $
+            application
+    application' request respond
+
+class HasEndpoint a where
+    getEndpoint :: Proxy a -> Request -> Maybe ([Text], Method)
+
+instance (HasEndpoint (a :: *), HasEndpoint (b :: *)) => HasEndpoint (a :<|> b) where
+    getEndpoint _ req =
+        getEndpoint (Proxy :: Proxy a) req <>
+        getEndpoint (Proxy :: Proxy b) req
+
+instance (KnownSymbol (path :: Symbol), HasEndpoint (sub :: *)) => HasEndpoint (path :> sub) where
+    getEndpoint _ req =
+        case pathInfo req of
+            p:ps | p == T.pack (symbolVal (Proxy :: Proxy path)) -> do
+                (end, method) <- getEndpoint (Proxy :: Proxy sub) req{ pathInfo = ps }
+                return (p:end, method)
+            _ -> Nothing
+
+instance (KnownSymbol (capture :: Symbol), HasEndpoint (sub :: *)) => HasEndpoint (Capture capture a :> sub) where
+    getEndpoint _ req =
+        case pathInfo req of
+            _:ps -> do
+                (end, method) <- getEndpoint (Proxy :: Proxy sub) req{ pathInfo = ps }
+                let p = T.pack $ (':':) $ symbolVal (Proxy :: Proxy capture)
+                return (p:end, method)
+            _ -> Nothing
+
+instance HasEndpoint (sub :: *) => HasEndpoint (Header h a :> sub) where
+    getEndpoint _ = getEndpoint (Proxy :: Proxy sub)
+
+instance HasEndpoint (sub :: *) => HasEndpoint (QueryParam (h :: Symbol) a :> sub) where
+    getEndpoint _ = getEndpoint (Proxy :: Proxy sub)
+
+instance HasEndpoint (sub :: *) => HasEndpoint (QueryParams (h :: Symbol) a :> sub) where
+    getEndpoint _ = getEndpoint (Proxy :: Proxy sub)
+
+instance HasEndpoint (sub :: *) => HasEndpoint (ReqBody a :> sub) where
+    getEndpoint _ = getEndpoint (Proxy :: Proxy sub)
+
+instance HasEndpoint (Get a) where
+    getEndpoint _ req = case pathInfo req of
+        [] | requestMethod req == "GET" -> Just ([],"GET")
+        _ -> Nothing
+
+instance HasEndpoint (Put a) where
+    getEndpoint _ req = case pathInfo req of
+        [] | requestMethod req == "PUT" -> Just ([],"PUT")
+        _ -> Nothing
+
+instance HasEndpoint (Post a) where
+    getEndpoint _ req = case pathInfo req of
+        [] | requestMethod req == "POST" -> Just ([],"POST")
+        _ -> Nothing
+
+instance HasEndpoint (Delete) where
+    getEndpoint _ req = case pathInfo req of
+        [] | requestMethod req == "DELETE" -> Just ([],"DELETE")
+        _ -> Nothing
+
+instance HasEndpoint (Raw) where
+    getEndpoint _ _ = Just ([],"RAW")

test/test.hs

@@ -0,0 +1,74 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE OverlappingInstances #-}
+
+import Control.Concurrent
+import Data.Aeson
+import Data.Monoid
+import Data.Proxy
+import Data.Text
+import GHC.Generics
+import Network.Wai
+import Network.Wai.Handler.Warp
+import System.Metrics
+import Servant
+
+import Servant.Ekg
+
+-- * Example
+
+-- | A greet message data type
+newtype Greet = Greet { _msg :: Text }
+  deriving (Generic, Show)
+
+instance FromJSON Greet
+instance ToJSON Greet
+
+-- API specification
+type TestApi =
+       -- GET /hello/:name?capital={true, false}  returns a Greet as JSON
+       "hello" :> Capture "name" Text :> QueryParam "capital" Bool :> Get Greet
+
+       -- POST /greet with a Greet as JSON in the request body,
+       --             returns a Greet as JSON
+  :<|> "greet" :> ReqBody Greet :> Post Greet
+
+       -- DELETE /greet/:greetid
+  :<|> "greet" :> Capture "greetid" Text :> Delete
+
+testApi :: Proxy TestApi
+testApi = Proxy
+
+-- Server-side handlers.
+--
+-- There's one handler per endpoint, which, just like in the type
+-- that represents the API, are glued together using :<|>.
+--
+-- Each handler runs in the 'EitherT (Int, String) IO' monad.
+server :: Server TestApi
+server = helloH :<|> postGreetH :<|> deleteGreetH
+
+  where helloH name Nothing = helloH name (Just False)
+        helloH name (Just False) = return . Greet $ "Hello, " <> name
+        helloH name (Just True) = return . Greet . toUpper $ "Hello, " <> name
+
+        postGreetH greet = return greet
+
+        deleteGreetH _ = return ()
+
+-- Turn the server into a WAI app. 'serve' is provided by servant,
+-- more precisely by the Servant.Server module.
+test :: Application
+test = serve testApi server
+
+-- Put this all to work!
+main :: IO ()
+main = do
+    ekg <- newStore
+    ms <- newMVar mempty
+    _ <- forkIO $ run 8001 $ monitorEndpoints testApi ekg ms test
+    return ()