feat: add metric pgrst_jwt_cache_size_bytes in admin server

Author: taimoorzaeem

CHANGELOG.md

@@ -17,6 +17,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
  - #2255, Apply `to_tsvector()` explicitly to the full-text search filtered column (excluding `tsvector` types) - @laurenceisla
  - #1578, Log the main SQL query to stderr at the current `log-level` when `log-query=main-query` - @laurenceisla
  - #3903, Log connection pool borrows on `log-level=debug` - @taimoorzaeem
+ - #3802, Add metric `pgrst_jwt_cache_size_bytes` in admin server - @taimoorzaeem
 
 ### Fixed
 

default.nix

@@ -85,6 +85,8 @@ rec {
     lib.enableExecutableProfiling
     lib.enableLibraryProfiling
     lib.dontHaddock
+    (drv: lib.appendConfigureFlags drv [ "--ghc-option=-DJWT_CACHE_METRIC" ])
+    (drv: lib.appendConfigureFlags drv [ "--flags=jwt-cache-metric" ])
   ];
 
   inherit (postgrest) env;

docs/references/observability.rst

@@ -195,6 +195,20 @@ pgrst_db_pool_max
 
 Max pool connections.
 
+JWT Cache Metric
+----------------
+
+Related to the :ref:`jwt_caching`.
+
+pgrst_jwt_cache_size_bytes
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+======== =======
+**Type** Gauge
+======== =======
+
+The JWT cache size in bytes.
+
 Traces
 ======
 

nix/overlays/haskell-packages.nix

@@ -50,6 +50,10 @@ let
       # jailbreak, because hspec limit for tests
       fuzzyset = prev.fuzzyset_0_2_4;
 
+      # TODO: Remove this once https://github.com/NixOS/nixpkgs/pull/375121
+      # has made it to us.
+      ghc-datasize = lib.markUnbroken prev.ghc-datasize;
+
       hasql-pool = lib.dontCheck (prev.callHackageDirect
         {
           pkg = "hasql-pool";

postgrest.cabal

@@ -40,6 +40,11 @@ flag hpc
   manual:      True
   description: Enable HPC (dev only)
 
+flag jwt-cache-metric
+  default:     False
+  manual:      True
+  description: Allow memory tests to run without ghc-datasize
+
 library
   default-language:   Haskell2010
   default-extensions: OverloadedStrings
@@ -110,6 +115,7 @@ library
                     , either                    >= 4.4.1 && < 5.1
                     , extra                     >= 1.7.0 && < 2.0
                     , fuzzyset                  >= 0.2.4 && < 0.3
+                    , ghc-heap                  >= 9.4 && < 9.9
                     , hasql                     >= 1.6.1.1 && < 1.7
                     , hasql-dynamic-statements  >= 0.3.1 && < 0.4
                     , hasql-notifications       >= 0.2.2.2 && < 0.2.3
@@ -168,6 +174,9 @@ library
   else
     ghc-options: -O2
 
+  if !flag(jwt-cache-metric)
+    build-depends: ghc-datasize >= 0.2.7 && < 0.3
+
   if !os(windows)
     build-depends:
       unix

src/PostgREST/Auth.hs

@@ -44,7 +44,7 @@ import System.IO.Unsafe        (unsafePerformIO)
 import System.TimeIt           (timeItT)
 
 import PostgREST.AppState      (AppState, getConfig, getJwtCacheState,
-                                getTime)
+                                getObserver, getTime)
 import PostgREST.Auth.JwtCache (lookupJwtCache)
 import PostgREST.Auth.Types    (AuthResult (..))
 import PostgREST.Config        (AppConfig (..), FilterExp (..),
@@ -161,6 +161,7 @@ middleware appState app req respond = do
   let token  = Wai.extractBearerAuth =<< lookup HTTP.hAuthorization (Wai.requestHeaders req)
       parseJwt = runExceptT $ parseToken conf token time >>= parseClaims conf
       jwtCacheState = getJwtCacheState appState
+      observer = getObserver appState
 
 -- If ServerTimingEnabled -> calculate JWT validation time
 -- If JwtCacheMaxLifetime -> cache JWT validation result
@@ -171,7 +172,7 @@ middleware appState app req respond = do
 
     (True, maxLifetime)  -> do
           (dur, authResult) <- timeItT $ case token of
-            Just tkn -> lookupJwtCache jwtCacheState tkn maxLifetime parseJwt time
+            Just tkn -> lookupJwtCache jwtCacheState tkn maxLifetime parseJwt time observer
             Nothing  -> parseJwt
           return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult & Vault.insert jwtDurKey dur }
 
@@ -181,7 +182,7 @@ middleware appState app req respond = do
 
     (False, maxLifetime) -> do
           authResult <- case token of
-            Just tkn -> lookupJwtCache jwtCacheState tkn maxLifetime parseJwt time
+            Just tkn -> lookupJwtCache jwtCacheState tkn maxLifetime parseJwt time observer
             Nothing -> parseJwt
           return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }
 

src/PostgREST/Auth/JwtCache.hs

@@ -4,7 +4,7 @@ Description : PostgREST Jwt Authentication Result Cache.
 
 This module provides functions to deal with the JWT cache
 -}
-{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE CPP #-}
 module PostgREST.Auth.JwtCache
   ( init
   , JwtCacheState
@@ -18,30 +18,42 @@ import qualified Data.Scientific   as Sci
 
 import Data.Time.Clock       (UTCTime, nominalDiffTimeToSeconds)
 import Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds)
-import System.Clock          (TimeSpec (..))
+#ifndef JWT_CACHE_METRIC
+import GHC.DataSize (recursiveSizeNF)
+#endif
+import System.Clock (TimeSpec (..))
 
-import PostgREST.Auth.Types (AuthResult (..))
-import PostgREST.Error      (Error (..))
+import PostgREST.Auth.Types  (AuthResult (..))
+import PostgREST.Error       (Error (..))
+import PostgREST.Observation (Observation (..), ObservationHandler)
 
+import Control.Debounce
 import Protolude
 
-newtype JwtCacheState = JwtCacheState
-  { jwtCache :: C.Cache ByteString AuthResult
+type SizeInBytes = Int
+
+-- TODO: EXPLAIN HERE WHY WE ARE STORING THE SIZE IN CACHE, Remove the below line
+-- https://github.com/PostgREST/postgrest/pull/3802#discussion_r1971074445
+data JwtCacheState = JwtCacheState
+  -- | Jwt Cache
+  { jwtCache                     :: C.Cache ByteString (AuthResult,SizeInBytes)
+  -- | Calculate cache size with debounce
+  , cacheSizeCalcDebounceTimeout :: MVar (IO ())
   }
 
 -- | Initialize JwtCacheState
 init :: IO JwtCacheState
 init = do
   cache <- C.newCache Nothing -- no default expiration
-  return $ JwtCacheState cache
+  JwtCacheState cache <$> newEmptyMVar
 
 -- | Used to retrieve and insert JWT to JWT Cache
-lookupJwtCache :: JwtCacheState -> ByteString -> Int -> IO (Either Error AuthResult) -> UTCTime -> IO (Either Error AuthResult)
-lookupJwtCache JwtCacheState{jwtCache} token maxLifetime parseJwt utc = do
-  checkCache <- C.lookup jwtCache token
-  authResult <- maybe parseJwt (pure . Right) checkCache
+lookupJwtCache :: JwtCacheState -> ByteString -> Int -> IO (Either Error AuthResult) -> UTCTime -> ObservationHandler -> IO (Either Error AuthResult)
+lookupJwtCache jwtCacheState token maxLifetime parseJwt utc observer = do
+  checkCache <- C.lookup (jwtCache jwtCacheState) token
+  authResult <- maybe parseJwt (pure . Right . fst) checkCache
 
-  case (authResult,checkCache) of
+  case (authResult, checkCache) of
     -- From comment:
     -- https://github.com/PostgREST/postgrest/pull/3801#discussion_r1857987914
     --
@@ -56,13 +68,20 @@ lookupJwtCache JwtCacheState{jwtCache} token maxLifetime parseJwt utc = do
 
     (Right res, Nothing) -> do -- cache miss
 
+      -- get expiration time
       let timeSpec = getTimeSpec res maxLifetime utc
 
       -- purge expired cache entries
-      C.purgeExpired jwtCache
+      C.purgeExpired (jwtCache jwtCacheState)
+
+      -- calculate size of the cache entry to store it with authResult
+      sz <- calcCacheEntrySizeInBytes (token,res,timeSpec)
 
-      -- insert new cache entry
-      C.insert' jwtCache (Just timeSpec) token res
+      -- insert new cache entry with byte size
+      C.insert' (jwtCache jwtCacheState) (Just timeSpec) token (res,sz)
+
+      -- calculate complete cache size with debounce and log it
+      updateCacheSizeWithDebounce jwtCacheState observer
 
     _                    -> pure ()
 
@@ -77,3 +96,61 @@ getTimeSpec res maxLifetime utc = do
   case expireJSON of
     Just (JSON.Number seconds) -> TimeSpec (sciToInt seconds - utcToSecs utc) 0
     _                          -> TimeSpec (fromIntegral maxLifetime :: Int64) 0
+
+-- | Update JwtCacheSize Metric
+--
+-- Runs the cache size calculation with debounce
+updateCacheSizeWithDebounce :: JwtCacheState -> ObservationHandler -> IO ()
+updateCacheSizeWithDebounce jwtCacheState observer = do
+  cSizeDebouncer <- tryReadMVar $ cacheSizeCalcDebounceTimeout jwtCacheState
+  case cSizeDebouncer of
+    Just d -> d
+    Nothing -> do
+      newDebouncer <-
+        mkDebounce defaultDebounceSettings
+          -- debounceFreq is set to default 1 second
+          { debounceAction = calculateSizeThenLog
+          , debounceEdge = leadingEdge -- logs at the start and the end
+          }
+      putMVar (cacheSizeCalcDebounceTimeout jwtCacheState) newDebouncer
+      newDebouncer
+    where
+      calculateSizeThenLog :: IO ()
+      calculateSizeThenLog = do
+        entries <- C.toList $ jwtCache jwtCacheState
+        -- extract the size from each entry and sum them all
+        let size = sum [ sz | (_,(_,sz),_) <- entries]
+        observer $ JwtCache size -- updates and logs the metric
+
+-- | Calculate JWT Cache Size in Bytes
+--
+-- The cache size is updated by calculating the size of every
+-- cache entry and updating the metric.
+--
+-- The cache entry consists of
+--   key          :: ByteString
+--   value        :: AuthReult
+--   expire value :: TimeSpec
+--
+-- We calculate the size of each cache entry component
+-- by using recursiveSizeNF function which first evaluates
+-- the data structure to Normal Form and then calculate size.
+-- The normal form evaluation is necessary for accurate size
+-- calculation because haskell is lazy and we dont wanna count
+-- the size of large thunks (unevaluated expressions)
+calcCacheEntrySizeInBytes :: (ByteString, AuthResult, TimeSpec) -> IO Int
+#ifndef JWT_CACHE_METRIC
+calcCacheEntrySizeInBytes entry = fromIntegral <$> getSize entry
+    where
+      -- We also include the size of SizeInBytes integer which is a constant 8 bytes
+      getSize :: (ByteString, AuthResult, TimeSpec) -> IO Word
+      getSize (bs, ar, ts) = do
+        keySize      <- recursiveSizeNF bs
+        arClaimsSize <- recursiveSizeNF $ authClaims ar
+        arRoleSize   <- recursiveSizeNF $ authRole ar
+        timeSpecSize <- liftA2 (+) (recursiveSizeNF (sec ts)) (recursiveSizeNF (nsec ts))
+        let sizeOfSizeEntryItself = 8 -- a constant 8 bytes size of each size entry in the cache
+        return (keySize + arClaimsSize + arRoleSize + timeSpecSize + sizeOfSizeEntryItself)
+#else
+calcCacheEntrySizeInBytes _ = return 0
+#endif

src/PostgREST/Logger.hs

@@ -100,6 +100,9 @@ observationLogger loggerState logLevel obs = case obs of
   o@PoolRequestFullfilled ->
     when (logLevel >= LogDebug) $ do
       logWithZTime loggerState $ observationMessage o
+  o@(JwtCache _) -> do
+    when (logLevel >= LogInfo) $ do
+      logWithZTime loggerState $ observationMessage o
   o ->
     logWithZTime loggerState $ observationMessage o
 

src/PostgREST/Metrics.hs

@@ -1,5 +1,5 @@
 {-|
-Module      : PostgREST.Logger
+Module      : PostgREST.Metrics
 Description : Metrics based on the Observation module. See Observation.hs.
 -}
 module PostgREST.Metrics
@@ -19,7 +19,7 @@ import PostgREST.Observation
 import Protolude
 
 data MetricsState =
-  MetricsState Counter Gauge Gauge Gauge (Vector Label1 Counter) Gauge
+  MetricsState Counter Gauge Gauge Gauge (Vector Label1 Counter) Gauge Gauge
 
 init :: Int -> IO MetricsState
 init configDbPoolSize = do
@@ -29,12 +29,13 @@ init configDbPoolSize = do
   poolMaxSize <- register $ gauge (Info "pgrst_db_pool_max" "Max pool connections")
   schemaCacheLoads <- register $ vector "status" $ counter (Info "pgrst_schema_cache_loads_total" "The total number of times the schema cache was loaded")
   schemaCacheQueryTime <- register $ gauge (Info "pgrst_schema_cache_query_time_seconds" "The query time in seconds of the last schema cache load")
+  jwtCacheSize <- register $ gauge (Info "pgrst_jwt_cache_size_bytes" "The JWT cache size in bytes")
   setGauge poolMaxSize (fromIntegral configDbPoolSize)
-  pure $ MetricsState poolTimeouts poolAvailable poolWaiting poolMaxSize schemaCacheLoads schemaCacheQueryTime
+  pure $ MetricsState poolTimeouts poolAvailable poolWaiting poolMaxSize schemaCacheLoads schemaCacheQueryTime jwtCacheSize
 
 -- Only some observations are used as metrics
 observationMetrics :: MetricsState -> ObservationHandler
-observationMetrics (MetricsState poolTimeouts poolAvailable poolWaiting _ schemaCacheLoads schemaCacheQueryTime) obs = case obs of
+observationMetrics (MetricsState poolTimeouts poolAvailable poolWaiting _ schemaCacheLoads schemaCacheQueryTime jwtCacheSize) obs = case obs of
   (PoolAcqTimeoutObs _) -> do
     incCounter poolTimeouts
   (HasqlPoolObs (SQL.ConnectionObservation _ status)) -> case status of
@@ -54,6 +55,8 @@ observationMetrics (MetricsState poolTimeouts poolAvailable poolWaiting _ schema
     setGauge schemaCacheQueryTime resTime
   SchemaCacheErrorObs _ -> do
     withLabel schemaCacheLoads "FAIL" incCounter
+  JwtCache cacheSize -> do
+    setGauge jwtCacheSize (fromIntegral cacheSize)
   _ ->
     pure ()
 

src/PostgREST/Observation.hs

@@ -59,6 +59,7 @@ data Observation
   | HasqlPoolObs SQL.Observation
   | PoolRequest
   | PoolRequestFullfilled
+  | JwtCache Int
 
 data ObsFatalError = ServerAuthError | ServerPgrstBug | ServerError42P05 | ServerError08P01
 
@@ -146,6 +147,8 @@ observationMessage = \case
     "Trying to borrow a connection from pool"
   PoolRequestFullfilled ->
     "Borrowed a connection from the pool"
+  JwtCache sz ->
+    "The JWT Cache size updated to " <> show sz <> " bytes"
   where
     showMillis :: Double -> Text
     showMillis x = toS $ showFFloat (Just 1) (x * 1000) ""