From 145c236a6250ccf0685ad423b03d28394499bdba Mon Sep 17 00:00:00 2001 From: "M. Taimoor Zaeem" Date: Sat, 23 Nov 2024 00:41:55 +0500 Subject: [PATCH] feat: add metric pgrst_jwt_cache_size_bytes in admin server --- docs/references/observability.rst | 14 ++++++++++ nix/overlays/haskell-packages.nix | 2 ++ postgrest.cabal | 2 ++ src/PostgREST/App.hs | 5 ++-- src/PostgREST/Auth.hs | 45 +++++++++++++++++++++++++------ src/PostgREST/Logger.hs | 3 +++ src/PostgREST/Metrics.hs | 11 +++++--- src/PostgREST/Observation.hs | 8 +++++- test/io/test_io.py | 2 ++ 9 files changed, 76 insertions(+), 16 deletions(-) diff --git a/docs/references/observability.rst b/docs/references/observability.rst index 04a261ba48..dcf7f02db4 100644 --- a/docs/references/observability.rst +++ b/docs/references/observability.rst @@ -169,6 +169,20 @@ pgrst_db_pool_max Max pool connections. +JWT Cache Metric +---------------- + +Related to the :ref:`jwt_caching`. + +pgrst_jwt_cache_size_bytes +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +======== ======= +**Type** Gauge +======== ======= + +Approximate JWT cache size in bytes. + Traces ====== diff --git a/nix/overlays/haskell-packages.nix b/nix/overlays/haskell-packages.nix index 917ab2478a..4da83005f5 100644 --- a/nix/overlays/haskell-packages.nix +++ b/nix/overlays/haskell-packages.nix @@ -69,6 +69,8 @@ let { postgresql = super.libpq; }); + + ghc-datasize = lib.markUnbroken prev.ghc-datasize; }; in { diff --git a/postgrest.cabal b/postgrest.cabal index 092a98de82..b204be4e7c 100644 --- a/postgrest.cabal +++ b/postgrest.cabal @@ -108,6 +108,8 @@ library , either >= 4.4.1 && < 5.1 , extra >= 1.7.0 && < 2.0 , fuzzyset >= 0.2.4 && < 0.3 + , ghc-datasize >= 0.2.7 && < 0.3 + , ghc-heap >= 9.4 && < 9.9 , gitrev >= 1.2 && < 1.4 , hasql >= 1.6.1.1 && < 1.7 , hasql-dynamic-statements >= 0.3.1 && < 0.4 diff --git a/src/PostgREST/App.hs b/src/PostgREST/App.hs index 99febebaca..3fd1f3c768 100644 --- a/src/PostgREST/App.hs +++ b/src/PostgREST/App.hs @@ -15,7 +15,7 @@ module PostgREST.App , run ) where - +import Control.Monad import Control.Monad.Except (liftEither) import Data.Either.Combinators (mapLeft) import Data.Maybe (fromJust) @@ -41,8 +41,7 @@ import qualified PostgREST.Response as Response import qualified PostgREST.Unix as Unix (installSignalHandlers) import PostgREST.ApiRequest (ApiRequest (..)) -import PostgREST.AppState (AppState) -import PostgREST.Auth (AuthResult (..)) +import PostgREST.AppState (AppState, AuthResult (..)) import PostgREST.Config (AppConfig (..), LogLevel (..)) import PostgREST.Config.PgVersion (PgVersion (..)) import PostgREST.Error (Error) diff --git a/src/PostgREST/Auth.hs b/src/PostgREST/Auth.hs index c8c3b720cb..4ea928a4dc 100644 --- a/src/PostgREST/Auth.hs +++ b/src/PostgREST/Auth.hs @@ -41,15 +41,17 @@ import Data.Either.Combinators (mapLeft) import Data.List (lookup) import Data.Time.Clock (UTCTime, nominalDiffTimeToSeconds) import Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds) +import GHC.DataSize (recursiveSizeNF) import System.Clock (TimeSpec (..)) import System.IO.Unsafe (unsafePerformIO) import System.TimeIt (timeItT) -import PostgREST.AppState (AppState, AuthResult (..), getConfig, - getJwtCache, getTime) -import PostgREST.Config (AppConfig (..), FilterExp (..), JSPath, - JSPathExp (..)) -import PostgREST.Error (Error (..)) +import PostgREST.AppState (AppState, AuthResult (..), getConfig, + getJwtCache, getObserver, getTime) +import PostgREST.Config (AppConfig (..), FilterExp (..), JSPath, + JSPathExp (..)) +import PostgREST.Error (Error (..)) +import PostgREST.Observation (Observation (..)) import Protolude @@ -153,7 +155,7 @@ middleware appState app req respond = do let token = fromMaybe "" $ Wai.extractBearerAuth =<< lookup HTTP.hAuthorization (Wai.requestHeaders req) parseJwt = runExceptT $ parseToken conf token time >>= parseClaims conf --- If DbPlanEnabled -> calculate JWT validation time +-- If ServerTimingEnabled -> calculate JWT validation time -- If JwtCacheMaxLifetime -> cache JWT validation result req' <- case (configServerTimingEnabled conf, configJwtCacheMaxLifetime conf) of (True, 0) -> do @@ -177,14 +179,24 @@ middleware appState app req respond = do -- | Used to retrieve and insert JWT to JWT Cache getJWTFromCache :: AppState -> ByteString -> Int -> IO (Either Error AuthResult) -> UTCTime -> IO (Either Error AuthResult) getJWTFromCache appState token maxLifetime parseJwt utc = do - checkCache <- C.lookup (getJwtCache appState) token + + checkCache <- C.lookup jwtCache token authResult <- maybe parseJwt (pure . Right) checkCache + -- if token not found, add to cache and increment cache size metric case (authResult,checkCache) of - (Right res, Nothing) -> C.insert' (getJwtCache appState) (getTimeSpec res maxLifetime utc) token res + (Right res, Nothing) -> do + let tSpec = getTimeSpec res maxLifetime utc + C.insert' jwtCache tSpec token res + entrySize <- calcCacheEntrySizeInBytes (token, res, tSpec) -- adds to cache + observer $ JWTCache entrySize + _ -> pure () return authResult + where + observer = getObserver appState + jwtCache = getJwtCache appState -- Used to extract JWT exp claim and add to JWT Cache getTimeSpec :: AuthResult -> Int -> UTCTime -> Maybe TimeSpec @@ -196,6 +208,23 @@ getTimeSpec res maxLifetime utc = do Just (JSON.Number seconds) -> Just $ TimeSpec (sciToInt seconds - utcToSecs utc) 0 _ -> Just $ TimeSpec (fromIntegral maxLifetime :: Int64) 0 +-- | Calculate a single entry of JWT Cache Size in Bytes +calcCacheEntrySizeInBytes :: (ByteString,AuthResult,Maybe TimeSpec) -> IO Int +calcCacheEntrySizeInBytes entry = do + sz <- getSize entry + return $ fromIntegral sz + where + getSize :: (ByteString, AuthResult,Maybe TimeSpec) -> IO Word + getSize (bs, ar, ts) = do + keySize <- recursiveSizeNF bs + arClaimsSize <- recursiveSizeNF $ authClaims ar + arRoleSize <- recursiveSizeNF $ authRole ar + timeSpecSize <- case ts of + Just TimeSpec{..} -> liftA2 (+) (recursiveSizeNF sec) (recursiveSizeNF nsec) + Nothing -> pure 0 + + return (keySize + arClaimsSize + arRoleSize + timeSpecSize) + authResultKey :: Vault.Key (Either Error AuthResult) authResultKey = unsafePerformIO Vault.newKey {-# NOINLINE authResultKey #-} diff --git a/src/PostgREST/Logger.hs b/src/PostgREST/Logger.hs index c224f74c79..dee387201f 100644 --- a/src/PostgREST/Logger.hs +++ b/src/PostgREST/Logger.hs @@ -88,6 +88,9 @@ observationLogger loggerState logLevel obs = case obs of o@(HasqlPoolObs _) -> do when (logLevel >= LogDebug) $ do logWithZTime loggerState $ observationMessage o + o@(JWTCache _) -> do + when (logLevel >= LogDebug) $ do + logWithZTime loggerState $ observationMessage o PoolRequest -> pure () PoolRequestFullfilled -> diff --git a/src/PostgREST/Metrics.hs b/src/PostgREST/Metrics.hs index 3999e43d83..8d2f74b50c 100644 --- a/src/PostgREST/Metrics.hs +++ b/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 entrySize -> do + addGauge jwtCacheSize (fromIntegral entrySize) _ -> pure () diff --git a/src/PostgREST/Observation.hs b/src/PostgREST/Observation.hs index 18fbf558d7..0846154b8c 100644 --- a/src/PostgREST/Observation.hs +++ b/src/PostgREST/Observation.hs @@ -57,8 +57,13 @@ data Observation | HasqlPoolObs SQL.Observation | PoolRequest | PoolRequestFullfilled + | JWTCache Int -data ObsFatalError = ServerAuthError | ServerPgrstBug | ServerError42P05 | ServerError08P01 +data ObsFatalError + = ServerAuthError + | ServerPgrstBug + | ServerError42P05 + | ServerError08P01 type ObservationHandler = Observation -> IO () @@ -138,6 +143,7 @@ observationMessage = \case SQL.ReleaseConnectionTerminationReason -> "release" SQL.NetworkErrorConnectionTerminationReason _ -> "network error" -- usage error is already logged, no need to repeat the same message. ) + JWTCache sz-> "The JWT Cache size increased to " <> show sz <> " bytes" _ -> mempty where showMillis :: Double -> Text diff --git a/test/io/test_io.py b/test/io/test_io.py index e61e074cbc..1d3a99b959 100644 --- a/test/io/test_io.py +++ b/test/io/test_io.py @@ -1632,6 +1632,8 @@ def test_admin_metrics(defaultenv): assert "pgrst_db_pool_available" in response.text assert "pgrst_db_pool_timeouts_total" in response.text + assert "pgrst_jwt_cache_size_bytes" in response.text + def test_schema_cache_startup_load_with_in_db_config(defaultenv, metapostgrest): "verify that the Schema Cache loads correctly at startup, using the in-db `pgrst.db_schemas` config"