simple prometheus metrics

1 files changed, 16 insertions, 5 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