Fix metrics server (parsonsmatt#35)

Author: parsonsmatt

src/Config.hs

@@ -2,9 +2,12 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE MultiParamTypeClasses      #-}
 {-# LANGUAGE OverloadedStrings          #-}
+
 {-# OPTIONS_GHC -fno-warn-orphans #-}
+
 module Config where
 
+import           Control.Concurrent                   (ThreadId)
 import           Control.Exception                    (throwIO)
 import           Control.Monad.Except                 (ExceptT, MonadError)
 import           Control.Monad.IO.Class
@@ -21,10 +24,10 @@ import           Database.Persist.Postgresql          (ConnectionPool,
                                                        ConnectionString,
                                                        createPostgresqlPool)
 import           Network.Wai                          (Middleware)
+import           Network.Wai.Handler.Warp             (Port)
 import           Network.Wai.Middleware.RequestLogger (logStdout, logStdoutDev)
 import           Servant                              (ServantErr)
 import           System.Environment                   (lookupEnv)
-import           Network.Wai.Handler.Warp             (Port)
 
 import           Logger
 
@@ -38,20 +41,23 @@ import           Logger
 newtype AppT m a
     = AppT
     { runApp :: ReaderT Config (ExceptT ServantErr m) a
-    } deriving ( Functor, Applicative, Monad, MonadReader Config,
-                 MonadError ServantErr, MonadIO)
+    } deriving
+    ( Functor, Applicative, Monad, MonadReader Config, MonadError ServantErr
+    , MonadIO
+    )
 
 type App = AppT IO
 
 -- | The Config for our application is (for now) the 'Environment' we're
 -- running in and a Persistent 'ConnectionPool'.
 data Config
     = Config
-    { configPool    :: ConnectionPool
-    , configEnv     :: Environment
-    , configMetrics :: Metrics
-    , configLogEnv  :: LogEnv
-    , configPort    :: Port
+    { configPool      :: ConnectionPool
+    , configEnv       :: Environment
+    , configMetrics   :: Metrics
+    , configEkgServer :: ThreadId
+    , configLogEnv    :: LogEnv
+    , configPort      :: Port
     }
 
 instance Monad m => MonadMetrics (AppT m) where

src/Init.hs

@@ -2,24 +2,27 @@
 
 module Init where
 
+import           Control.Concurrent          (killThread)
 import qualified Control.Monad.Metrics       as M
 import           Database.Persist.Postgresql (runSqlPool)
+import           Lens.Micro                  ((^.))
 import           Network.Wai                 (Application)
 import           Network.Wai.Metrics         (metrics, registerWaiMetrics)
 import           System.Environment          (lookupEnv)
-import           System.Remote.Monitoring    (forkServer, serverMetricStore)
+import           System.Remote.Monitoring    (forkServer, serverMetricStore,
+                                              serverThreadId)
 
 import           Api                         (app)
 import           Api.User                    (generateJavaScript)
 import           Config                      (Config (..), Environment (..),
                                               makePool, setLogger)
+import           Control.Exception           (bracket)
+import qualified Data.Pool                   as Pool
+import qualified Katip
 import           Logger                      (defaultLogEnv)
 import           Models                      (doMigrations)
-import           Safe                        (readMay)
-import           Control.Exception           (bracket)
 import           Network.Wai.Handler.Warp    (run)
-import qualified Data.Pool as Pool
-import qualified Katip
+import           Safe                        (readMay)
 
 -- | An action that creates a WAI 'Application' together with its resources,
 --   runs it, and tears it down on exit
@@ -31,14 +34,12 @@ runApp = bracket acquireConfig shutdownApp runApp
 -- | The 'initialize' function accepts the required environment information,
 -- initializes the WAI 'Application' and returns it
 initialize :: Config -> IO Application
-initialize cfg@(Config pool env _ _ _) = do
-    waiMetrics <-
-        registerWaiMetrics =<< serverMetricStore
-        <$> forkServer "localhost" 8000
-    let logger = setLogger env
-    runSqlPool doMigrations pool
+initialize cfg = do
+    waiMetrics <- registerWaiMetrics (configMetrics cfg ^. M.metricsStore)
+    let logger = setLogger (configEnv cfg)
+    runSqlPool doMigrations (configPool cfg)
     generateJavaScript
-    pure (logger $ metrics waiMetrics $ app cfg)
+    pure . logger . metrics waiMetrics . app $ cfg
 
 -- | Allocates resources for 'Config'
 acquireConfig :: IO Config
@@ -47,23 +48,28 @@ acquireConfig = do
     env  <- lookupSetting "ENV" Development
     logEnv <- defaultLogEnv
     pool <- makePool env logEnv
-    store <- serverMetricStore <$> forkServer "localhost" 8000
-    waiMetrics <- registerWaiMetrics store
+    ekgServer <- forkServer "localhost" 8000
+    let store = serverMetricStore ekgServer
+    waiMetrics <- registerWaiMetrics  store
     metr <- M.initializeWith store
-    pure Config { configPool = pool
-                , configEnv = env
-                , configMetrics = metr
-                , configLogEnv = logEnv
-                , configPort = port }
+    pure Config
+        { configPool = pool
+        , configEnv = env
+        , configMetrics = metr
+        , configLogEnv = logEnv
+        , configPort = port
+        , configEkgServer = serverThreadId ekgServer
+        }
 
 -- | Takes care of cleaning up 'Config' resources
 shutdownApp :: Config -> IO ()
-shutdownApp (Config pool _env metr logEnv _port) = do
-    Katip.closeScribes logEnv
-    Pool.destroyAllResources pool
+shutdownApp cfg = do
+    Katip.closeScribes (configLogEnv cfg)
+    Pool.destroyAllResources (configPool cfg)
     -- Monad.Metrics does not provide a function to destroy metrics store
     -- so, it'll hopefully get torn down when async exception gets thrown
     -- at metrics server process
+    killThread (configEkgServer cfg)
     pure ()
 
 -- | Looks up a setting in the environment, with a provided default, and