simple prometheus metrics

2 files changed, 24 insertions, 6 deletions

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] }