refactor(test): provide means to validate metrics and observations

Some helpers are provided for introspecting metrics already (used in JWT cache tests). This change provides facilities to additionally validate emited Observation events.
A new Spec module is also implemented, adding basic tests of schema cache reloading - their main goal is to excercise the new infrastructure.

Author: mkleczek

postgrest.cabal

@@ -219,6 +219,7 @@ test-suite spec
                       Feature.ConcurrentSpec
                       Feature.CorsSpec
                       Feature.ExtraSearchPathSpec
+                      Feature.MetricsSpec
                       Feature.NoSuperuserSpec
                       Feature.ObservabilitySpec
                       Feature.OpenApi.DisabledOpenApiSpec

src/PostgREST/AppState.hs

@@ -18,6 +18,7 @@ module PostgREST.AppState
   , init
   , initSockets
   , initWithPool
+  , putConfig -- For tests TODO refactoring
   , putNextListenerDelay
   , putSchemaCache
   , putPgVersion

test/spec/Feature/MetricsSpec.hs

@@ -0,0 +1,49 @@
+{-# LANGUAGE DataKinds           #-}
+{-# LANGUAGE FlexibleContexts    #-}
+{-# LANGUAGE MonadComprehensions #-}
+{-# LANGUAGE TypeApplications    #-}
+
+module Feature.MetricsSpec where
+
+import           Network.Wai           (Application)
+import qualified PostgREST.AppState    as AppState
+import           PostgREST.Config      (AppConfig (configDbSchemas))
+import qualified PostgREST.Metrics     as Metrics
+import           PostgREST.Observation
+import           Prometheus            (getCounter, getVectorWith)
+import           Protolude
+import           SpecHelper
+import           Test.Hspec            (SpecWith, describe, it)
+
+spec :: SpecWith (((AppState.AppState, Metrics.MetricsState), Chan Observation), Application)
+spec = describe "Server started with metrics enabled" $ do
+  it "Should update pgrst_schema_cache_loads_total[SUCCESS]" $ do
+    ((appState, metrics), waitFor) <- prepareState
+
+    liftIO $ checkState' metrics [
+        schemaCacheLoads "SUCCESS" (+1)
+      ] $ do
+        AppState.schemaCacheLoader appState
+        waitFor (1 * sec) "SchemaCacheLoadedObs" $ \x -> [ o | o@(SchemaCacheLoadedObs{}) <- pure x]
+
+  it "Should update pgrst_schema_cache_loads_total[ERROR]" $ do
+    ((appState, metrics), waitFor) <- prepareState
+
+    liftIO $ checkState' metrics [
+        schemaCacheLoads "FAIL" (+1),
+        schemaCacheLoads "SUCCESS" (+1)
+      ] $ do
+        AppState.getConfig appState >>= \prev -> do
+          AppState.putConfig appState $ prev { configDbSchemas = pure "bad_schema" }
+          AppState.schemaCacheLoader appState
+          waitFor (1 * sec) "SchemaCacheErrorObs" $ \x -> [ o | o@(SchemaCacheErrorObs{}) <- pure x]
+          AppState.putConfig appState prev
+
+        -- wait up to 2 secs so that retry can happen
+        waitFor (2 * sec) "SchemaCacheLoadedObs" $ \x -> [ o | o@(SchemaCacheLoadedObs{}) <- pure x]
+
+  where
+    -- prometheus-client api to handle vectors is convoluted
+    schemaCacheLoads label = expectField @"schemaCacheLoads" $
+      (foldMap (Sum . round @Double @Int . snd) . find ((== label) . fst) <$>) . (`getVectorWith` getCounter)
+    sec = 1000000

test/spec/Main.hs

@@ -29,6 +29,7 @@ import qualified Feature.Auth.NoJwtSecretSpec
 import qualified Feature.ConcurrentSpec
 import qualified Feature.CorsSpec
 import qualified Feature.ExtraSearchPathSpec
+import qualified Feature.MetricsSpec
 import qualified Feature.NoSuperuserSpec
 import qualified Feature.ObservabilitySpec
 import qualified Feature.OpenApi.DisabledOpenApiSpec
@@ -68,16 +69,23 @@ import qualified Feature.Query.UpdateSpec
 import qualified Feature.Query.UpsertSpec
 import qualified Feature.RollbackSpec
 import qualified Feature.RpcPreRequestGucsSpec
+import           PostgREST.Observation                 (Observation (HasqlPoolObs))
 
 
 main :: IO ()
 main = do
+  poolChan <- newChan
+  -- make sure poolChan is not growing indefinitely
+  void $ forkIO $ fix (readChan poolChan *>)
+  metricsState <- Metrics.init (configDbPoolSize testCfg)
   pool <- P.acquire $ P.settings
     [ P.size 3
     , P.acquisitionTimeout 10
     , P.agingTimeout 60
     , P.idlenessTimeout 60
     , P.staticConnectionSettings (toUtf8 $ configDbUri testCfg)
+    -- make sure metrics are updated and pool observations published to poolChan
+    , P.observationHandler $ (writeChan poolChan <> Metrics.observationMetrics metricsState) . HasqlPoolObs
     ]
 
   actualPgVersion <- either (panic . show) id <$> P.use pool (queryPgVersion False)
@@ -86,7 +94,6 @@ main = do
   baseSchemaCache <- loadSCache pool testCfg
   sockets <- AppState.initSockets testCfg
   loggerState <- Logger.init
-  metricsState <- Metrics.init (configDbPoolSize testCfg)
 
   let
     initApp sCache st config = do
@@ -95,6 +102,14 @@ main = do
       AppState.putSchemaCache appState (Just sCache)
       return (st, postgrest (configLogLevel config) appState (pure ()))
 
+    initObservationsApp sCache config = do
+      -- duplicate poolChan as a starting point
+      obsChan <- dupChan poolChan
+      appState <- AppState.initWithPool sockets pool config loggerState metricsState (Metrics.observationMetrics metricsState <> writeChan obsChan)
+      AppState.putPgVersion appState actualPgVersion
+      AppState.putSchemaCache appState (Just sCache)
+      return (((appState, metricsState), obsChan), postgrest (configLogLevel config) appState (pure ()))
+
     -- For tests that run with the same schema cache
     app = initApp baseSchemaCache ()
 
@@ -123,6 +138,7 @@ main = do
       obsApp               = app testObservabilityCfg
       serverTiming         = app testCfgServerTiming
       aggregatesEnabled    = app testCfgAggregatesEnabled
+      observationsApp      = initObservationsApp baseSchemaCache testCfg
 
       extraSearchPathApp   = appDbs testCfgExtraSearchPath
       unicodeApp           = appDbs testUnicodeCfg
@@ -278,6 +294,12 @@ main = do
     before (initApp baseSchemaCache metricsState testCfgJwtCache) $
       describe "Feature.Auth.JwtCacheSpec" Feature.Auth.JwtCacheSpec.spec
 
+    before (initApp baseSchemaCache metricsState testCfgJwtCache) $
+      describe "Feature.Auth.JwtCacheSpec" Feature.Auth.JwtCacheSpec.spec
+
+    before observationsApp $
+      describe "Feature.MetricsSpec" Feature.MetricsSpec.spec
+
   where
     loadSCache pool conf =
       either (panic.show) id <$> P.use pool (HT.transaction HT.ReadCommitted HT.Read $ querySchemaCache conf)

test/spec/SpecHelper.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE AllowAmbiguousTypes       #-}
+{-# LANGUAGE DeriveAnyClass            #-}
 {-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE FlexibleContexts          #-}
 {-# LANGUAGE RankNTypes                #-}
@@ -41,10 +42,13 @@ import PostgREST.Config                  (AppConfig (..),
                                           LogLevel (..),
                                           OpenAPIMode (..),
                                           parseSecret)
+import PostgREST.Observation             (Observation,
+                                          observationMessage)
 import PostgREST.SchemaCache.Identifiers (QualifiedIdentifier (..))
 import Prometheus                        (Counter, getCounter)
 import Protolude                         hiding (get, toS)
 import Protolude.Conv                    (toS)
+import System.Timeout                    (timeout)
 import Test.Hspec.Expectations.Contrib   (annotate)
 
 filterAndMatchCT :: BS.ByteString -> MatchHeader
@@ -380,3 +384,45 @@ expectCounter :: forall s st m. (KnownSymbol s, HasField s st Counter, MonadIO m
 expectCounter = expectField @s intCounter
   where
     intCounter = ((round @Double @Int) <$>) . getCounter
+
+data TimeoutException = TimeoutException deriving (Show, Exception)
+
+accumulateUntilTimeout :: Int -> (s -> a -> s) -> s -> IO a -> IO s
+accumulateUntilTimeout t f start act = do
+  tid <- myThreadId
+  -- mask to make sure TimeoutException is not thrown before starting the loop
+  mask $ \unmask -> do
+    -- start timeout thread unmasking exceptions
+    ttid <- forkIOWithUnmask ($ (threadDelay t *> throwTo tid TimeoutException))
+    -- unmask effect
+    unmask $ fix (\loop accum -> (act >>= loop . f accum) `onTimeout` pure accum) start
+      -- make sure we catch timeout if happens bifore entering the loop
+      `onTimeout` pure start
+      -- make sure timer thread is killed on other exceptions
+      -- so that it won't throw TimeoutException later
+      `onException` killThread ttid
+  where
+    onTimeout m a = m `catch` \TimeoutException -> a
+
+
+prepareState :: HasCallStack => Traversable f => WaiSession (f (Chan Observation)) (f (Int -> String -> (Observation -> Maybe a) -> IO ()))
+prepareState = getState >>= traverse (liftA2 (<$>) waitFor (liftIO . dupChan))
+  where
+    -- read messages from copy chan and once condition is met drain original to the same point
+    -- upon timeout report error and messages remaining in the original chan
+    -- that way we report messages since last successful read
+    waitFor orig copy t msg f =
+      timeout t (readUntil copy *> readUntil orig) >>= maybe failTimeout mempty
+      where
+        failTimeout =
+          takeUntilTimeout 100000 (readChan orig)
+            >>= expectationFailure
+            . ("Timeout waiting for " <> msg <> " at " <> loc <> ". Remaining observations:\n" ++)
+            . foldMap ((++ "\n") . show . observationMessage)
+        readUntil = void . untilM (pure . isJust . f) . readChan
+    loc = fold (head (prettySrcLoc . snd <$> getCallStack callStack))
+    -- execute effectful computation until result meets provided condition
+    untilM cond m = fix $ \loop -> m >>= \a -> ifM (cond a) (pure a) loop
+    -- duplicate the provided channel and construct wairFor function binding both channels
+    -- accumulate effecful computation results into a list for specified time
+    takeUntilTimeout t = fmap reverse . accumulateUntilTimeout t (flip (:)) []