From 7ff02dbf70b9a0f7a12cc4451de2735946252e60 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 --- .github/workflows/test.yaml | 2 +- docs/references/observability.rst | 14 ++++++++++ nix/overlays/haskell-packages.nix | 5 ++++ 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 | 2 ++ stack-21.7.yaml | 1 + stack.yaml | 1 + test/io/test_io.py | 2 ++ 12 files changed, 77 insertions(+), 16 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 7cfedb0b2b..34012a454d 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -128,5 +128,5 @@ jobs: prefix: v - name: Run loadtest run: | - postgrest-loadtest-against main ${{ steps.get-latest-tag.outputs.tag }} + PGRST_BUILD_CABAL=1 postgrest-loadtest-against main ${{ steps.get-latest-tag.outputs.tag }} postgrest-loadtest-report >> "$GITHUB_STEP_SUMMARY" 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..043f333c6e 100644 --- a/nix/overlays/haskell-packages.nix +++ b/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"; @@ -69,6 +73,7 @@ let { postgresql = super.libpq; }); + }; 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..a40c9c9181 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..3cb1779041 100644 --- a/src/PostgREST/Observation.hs +++ b/src/PostgREST/Observation.hs @@ -57,6 +57,7 @@ data Observation | HasqlPoolObs SQL.Observation | PoolRequest | PoolRequestFullfilled + | JWTCache Int data ObsFatalError = ServerAuthError | ServerPgrstBug | ServerError42P05 | ServerError08P01 @@ -138,6 +139,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/stack-21.7.yaml b/stack-21.7.yaml index e75e6a77cb..21c242aa6c 100644 --- a/stack-21.7.yaml +++ b/stack-21.7.yaml @@ -19,6 +19,7 @@ nix: extra-deps: - configurator-pg-0.2.10 - fuzzyset-0.2.4 + - ghc-datasize-0.2.7 - hasql-notifications-0.2.2.0 - hasql-pool-1.0.1 - postgresql-libpq-0.10.1.0 diff --git a/stack.yaml b/stack.yaml index 477fe59f31..f9cf8a3bc1 100644 --- a/stack.yaml +++ b/stack.yaml @@ -11,6 +11,7 @@ nix: extra-deps: - fuzzyset-0.2.4 + - ghc-datasize-0.2.7 - hasql-pool-1.0.1 - jose-jwt-0.10.0 - postgresql-libpq-0.10.1.0 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"