simple prometheus metrics

3 files changed, 35 insertions, 7 deletions

diff --git a/lib/API.hs b/lib/API.hs
index 32465c7..70971c3 100644
--- a/lib/API.hs
+++ b/lib/API.hs
@@ -5,7 +5,7 @@
 
 -- | The sole authorative definition of this server's API, given as a Servant-style
 -- Haskell type. All other descriptions of the API are generated from this one.
-module API (API, CompleteAPI, GtfsRealtimeAPI, RegisterJson(..)) where
+module API (API, CompleteAPI, GtfsRealtimeAPI, RegisterJson(..), Metrics(..)) where
 
 import           Data.Map                    (Map)
 import           Data.Proxy                  (Proxy (..))
@@ -32,11 +32,14 @@ import           Web.Internal.FormUrlEncoded (Form)
 
 import           Control.Lens                (At (at), (&), (?~))
 import           Data.Aeson                  (FromJSON (..), genericParseJSON)
+import           Data.ByteString.Lazy        (ByteString)
 import           Data.HashMap.Strict.InsOrd  (singleton)
 import           GHC.Generics                (Generic)
 import           GTFS
 import           GTFS.Realtime.FeedEntity
 import           GTFS.Realtime.FeedMessage   (FeedMessage)
+import           Prometheus
+
 import           Persist
 
 newtype RegisterJson = RegisterJson
@@ -76,9 +79,16 @@ type GtfsRealtimeAPI = "servicealerts" :> Get '[Proto] FeedMessage
 type CompleteAPI =
   "api" :> "openapi" :> Get '[JSON] Swagger
   :<|> "api" :> API
+  :<|> "metrics" :> Get '[PlainText] Text
   :<|> Raw -- hook for yesod frontend
 
 
+data Metrics = Metrics
+  { metricsWSGauge :: Gauge }
+
+
+
+
 -- TODO write something useful here! (and if it's just "hey this is some websocket thingie")
 instance HasSwagger WebSocket where
   toSwagger _ = mempty
diff --git a/lib/Server.hs b/lib/Server.hs
index 93046f8..6ca9c14 100644
--- a/lib/Server.hs
+++ b/lib/Server.hs
@@ -12,6 +12,7 @@
 -- Implementation of the API. This module is the main point of the program.
 module Server (application) where
 import           Control.Monad               (forever, unless, void, when)
+import           Control.Monad.Catch         (handle)
 import           Control.Monad.Extra         (ifM, maybeM, unlessM, whenM)
 import           Control.Monad.IO.Class      (MonadIO (liftIO))
 import           Control.Monad.Logger        (LoggingT, logWarnN)
@@ -26,6 +27,7 @@ import           Data.Pool                   (Pool)
 import           Data.Proxy                  (Proxy (Proxy))
 import           Data.Swagger                (toSchema)
 import           Data.Text                   (Text)
+import           Data.Text.Encoding          (decodeUtf8)
 import           Data.Time                   (NominalDiffTime,
                                               UTCTime (utctDay), addUTCTime,
                                               diffUTCTime, getCurrentTime,
@@ -58,11 +60,17 @@ import           System.IO.Unsafe
 
 import           Config                      (ServerConfig)
 import           Data.ByteString             (ByteString)
+import           Data.ByteString.Lazy        (toStrict)
+import           Prometheus
+import           Prometheus.Metric.GHC
 
 application :: GTFS -> Pool SqlBackend -> IO Application
 application gtfs dbpool = do
   doMigration dbpool
-  pure $ serve (Proxy @CompleteAPI) $ hoistServer (Proxy @CompleteAPI) runService $ server gtfs dbpool
+  metrics <- Metrics
+    <$> register (gauge (Info "ws_connections" "Number of WS Connections"))
+  register ghcMetrics
+  pure $ serve (Proxy @CompleteAPI) $ hoistServer (Proxy @CompleteAPI) runService $ server gtfs metrics dbpool
 
 -- databaseMigration :: ConnectionString -> IO ()
 doMigration pool = runSql pool $
@@ -71,12 +79,12 @@ doMigration pool = runSql pool $
   -- returns an empty list
   runMigration migrateAll
 
-server :: GTFS -> Pool SqlBackend -> Service CompleteAPI
-server gtfs@GTFS{..} dbpool = handleDebugAPI
+server :: GTFS -> Metrics -> Pool SqlBackend -> Service CompleteAPI
+server gtfs@GTFS{..} Metrics{..} dbpool = handleDebugAPI
   :<|> (handleStations :<|> handleTimetable :<|> handleTrip
         :<|> handleRegister :<|> handleTrainPing (throwError err401) :<|> handleWS
         :<|> handleDebugState :<|> handleDebugTrain :<|> handleDebugRegister
-        :<|> gtfsRealtimeServer gtfs dbpool)
+        :<|> gtfsRealtimeServer gtfs dbpool) :<|> metrics
   :<|> pure (unsafePerformIO (toWaiAppPlain (ControlRoom gtfs dbpool)))
   where handleStations = pure stations
         handleTimetable station maybeDay = do
@@ -118,12 +126,14 @@ server gtfs@GTFS{..} dbpool = handleDebugAPI
             pure (Just anchor)
         handleWS conn = do
           liftIO $ WS.forkPingThread conn 30
-          forever $ do
+          incGauge metricsWSGauge
+          handle (\(e :: WS.ConnectionException) -> decGauge metricsWSGauge) $ forever $ do
             msg <- liftIO $ WS.receiveData conn
             case A.eitherDecode msg of
               Left err -> do
                 logWarnN ("stray websocket message: "+|show msg|+" (could not decode: "+|err|+")")
                 liftIO $ WS.sendClose conn (C8.pack err)
+                decGauge metricsWSGauge
               Right ping ->
                 -- if invalid token, send a "polite" close request. Note that the client may
                 -- ignore this and continue sending messages, which will continue to be handled.
@@ -147,6 +157,7 @@ server gtfs@GTFS{..} dbpool = handleDebugAPI
               selectList [TrainPingToken ==. token] [] <&> fmap entityVal
             pure (concat pings)
         handleDebugAPI = pure $ toSwagger (Proxy @API)
+        metrics = exportMetricsAsText <&> (decodeUtf8 . toStrict)
 
 
 -- TODO: proper debug logging for expired tokens
diff --git a/lib/Server/Util.hs b/lib/Server/Util.hs
index 4410711..41d26f7 100644
--- a/lib/Server/Util.hs
+++ b/lib/Server/Util.hs
@@ -1,4 +1,6 @@
-{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleContexts     #-}
+{-# LANGUAGE FlexibleInstances    #-}
+{-# LANGUAGE TypeSynonymInstances #-}
 
 -- | mostly the monad the service runs in
 module Server.Util (Service, ServiceM, runService, sendErrorMsg, secondsNow, utcToSeconds) where
@@ -12,6 +14,7 @@ import           Data.Time              (Day, UTCTime (..), diffUTCTime,
                                          getCurrentTime,
                                          nominalDiffTimeToSeconds)
 import           GTFS                   (Seconds (..))
+import           Prometheus             (MonadMonitor (doIO))
 import           Servant                (Handler, ServerError, ServerT, err404,
                                          errBody, errHeaders, throwError)
 
@@ -21,6 +24,10 @@ type Service api = ServerT api ServiceM
 runService :: ServiceM a -> Handler a
 runService = runStderrLoggingT
 
+instance MonadMonitor ServiceM where
+  doIO = liftIO
+
+
 sendErrorMsg :: Text -> ServiceM a
 sendErrorMsg msg = throwError err404
   { errBody = A.encode $ A.object ["error" A..= (404 :: Int), "msg" A..= msg] }