feat: add more perf counters to Server-Timing (#2983)

2023-10-12 09:55:30 -03:00
parent 056c748c5f
commit dc01c748ae
9 changed files with 180 additions and 129 deletions
@@ -45,23 +45,27 @@ import qualified PostgREST.Plan             as Plan
 import qualified PostgREST.Query            as Query
 import qualified PostgREST.Response         as Response

-import PostgREST.ApiRequest          (Action (..), ApiRequest (..),
-                                      Mutation (..), Target (..))
-import PostgREST.AppState            (AppState)
-import PostgREST.Auth                (AuthResult (..))
-import PostgREST.Config              (AppConfig (..))
-import PostgREST.Config.PgVersion    (PgVersion (..))
-import PostgREST.Error               (Error)
-import PostgREST.Query               (DbHandler)
-import PostgREST.Response            (ServerTimingParams (..))
-import PostgREST.SchemaCache         (SchemaCache (..))
-import PostgREST.SchemaCache.Routine (Routine (..))
-import PostgREST.Version             (docsVersion, prettyVersion)
+import PostgREST.ApiRequest           (Action (..), ApiRequest (..),
+                                       Mutation (..), Target (..))
+import PostgREST.AppState             (AppState)
+import PostgREST.Auth                 (AuthResult (..))
+import PostgREST.Config               (AppConfig (..))
+import PostgREST.Config.PgVersion     (PgVersion (..))
+import PostgREST.Error                (Error)
+import PostgREST.Query                (DbHandler)
+import PostgREST.Response.Performance (ServerMetric (..),
+                                       ServerTimingData,
+                                       renderServerTimingHeader)
+import PostgREST.SchemaCache          (SchemaCache (..))
+import PostgREST.SchemaCache.Routine  (Routine (..))
+import PostgREST.Version              (docsVersion, prettyVersion)

 import qualified Data.ByteString.Char8 as BS
 import qualified Data.List             as L
+import qualified Data.Map              as Map (fromList)
 import qualified Network.HTTP.Types    as HTTP
 import           Protolude             hiding (Handler)
+import           System.TimeIt         (timeItT)

 type Handler = ExceptT Error

@@ -155,8 +159,8 @@ postgrestResponse appState conf@AppConfig{..} maybeSchemaCache pgVer authResult@
    liftEither . mapLeft Error.ApiRequestError $
      ApiRequest.userApiRequest conf req body

-  let serverTimingParams = if configDbPlanEnabled then Just (ServerTimingParams { jwtDur = fromJust $ Auth.getJwtDur req }) else Nothing
-  handleRequest authResult conf appState (Just authRole /= configDbAnonRole) configDbPreparedStatements pgVer apiRequest sCache serverTimingParams
+  let jwtTiming = (SMJwt, if configDbPlanEnabled then Auth.getJwtDur req else Nothing)
+  handleRequest authResult conf appState (Just authRole /= configDbAnonRole) configDbPreparedStatements pgVer apiRequest sCache jwtTiming

 runDbHandler :: AppState.AppState -> SQL.IsolationLevel -> SQL.Mode -> Bool -> Bool -> DbHandler b -> Handler IO b
 runDbHandler appState isoLvl mode authenticated prepared handler = do
@@ -170,63 +174,73 @@ runDbHandler appState isoLvl mode authenticated prepared handler = do

  liftEither resp

-handleRequest :: AuthResult -> AppConfig -> AppState.AppState -> Bool -> Bool -> PgVersion -> ApiRequest -> SchemaCache -> Maybe ServerTimingParams -> Handler IO Wai.Response
-handleRequest AuthResult{..} conf appState authenticated prepared pgVer apiReq@ApiRequest{..} sCache serverTimingParams =
+handleRequest :: AuthResult -> AppConfig -> AppState.AppState -> Bool -> Bool -> PgVersion -> ApiRequest -> SchemaCache -> (ServerMetric, Maybe Double) -> Handler IO Wai.Response
+handleRequest AuthResult{..} conf appState authenticated prepared pgVer apiReq@ApiRequest{..} sCache jwtTime =
  case (iAction, iTarget) of
    (ActionRead headersOnly, TargetIdent identifier) -> do
-      wrPlan <- liftEither $ Plan.wrappedReadPlan identifier conf sCache apiReq
-      resultSet <- runQuery roleIsoLvl (Plan.wrTxMode wrPlan) $ Query.readQuery wrPlan conf apiReq
-      pgrst <- liftEither $ Response.readResponse wrPlan headersOnly identifier apiReq resultSet serverTimingParams
-      return $ pgrstResponse pgrst
+      (planTime', wrPlan) <- withTiming $ liftEither $ Plan.wrappedReadPlan identifier conf sCache apiReq
+      (rsTime', resultSet) <- withTiming $ runQuery roleIsoLvl (Plan.wrTxMode wrPlan) $ Query.readQuery wrPlan conf apiReq
+      (renderTime', pgrst) <- withTiming $ liftEither $ Response.readResponse wrPlan headersOnly identifier apiReq resultSet
+      let metrics = Map.fromList [(SMPlan, planTime'), (SMQuery, rsTime'), (SMRender, renderTime'), jwtTime]
+      return $ pgrstResponse metrics pgrst

    (ActionMutate MutationCreate, TargetIdent identifier) -> do
-      mrPlan <- liftEither $ Plan.mutateReadPlan MutationCreate apiReq identifier conf sCache
-      resultSet <- runQuery roleIsoLvl (Plan.mrTxMode mrPlan) $ Query.createQuery mrPlan apiReq conf
-      pgrst <- liftEither $ Response.createResponse identifier mrPlan apiReq resultSet serverTimingParams
-      return $ pgrstResponse pgrst
+      (planTime', mrPlan) <- withTiming $ liftEither $ Plan.mutateReadPlan MutationCreate apiReq identifier conf sCache
+      (rsTime', resultSet) <- withTiming $ runQuery roleIsoLvl (Plan.mrTxMode mrPlan) $ Query.createQuery mrPlan apiReq conf
+      (renderTime', pgrst) <- withTiming $ liftEither $ Response.createResponse identifier mrPlan apiReq resultSet
+      let metrics = Map.fromList [(SMPlan, planTime'), (SMQuery, rsTime'), (SMRender, renderTime'), jwtTime]
+      return $ pgrstResponse metrics pgrst

    (ActionMutate MutationUpdate, TargetIdent identifier) -> do
-      mrPlan <- liftEither $ Plan.mutateReadPlan MutationUpdate apiReq identifier conf sCache
-      resultSet <- runQuery roleIsoLvl (Plan.mrTxMode mrPlan) $ Query.updateQuery mrPlan apiReq conf
-      pgrst <- liftEither $ Response.updateResponse mrPlan apiReq resultSet serverTimingParams
-      return $ pgrstResponse pgrst
+      (planTime', mrPlan) <- withTiming $ liftEither $ Plan.mutateReadPlan MutationUpdate apiReq identifier conf sCache
+      (rsTime', resultSet) <- withTiming $ runQuery roleIsoLvl (Plan.mrTxMode mrPlan) $ Query.updateQuery mrPlan apiReq conf
+      (renderTime', pgrst) <- withTiming $ liftEither $ Response.updateResponse mrPlan apiReq resultSet
+      let metrics = Map.fromList [(SMPlan, planTime'), (SMQuery, rsTime'), (SMRender, renderTime'), jwtTime]
+      return $ pgrstResponse metrics pgrst

    (ActionMutate MutationSingleUpsert, TargetIdent identifier) -> do
-      mrPlan <- liftEither $ Plan.mutateReadPlan MutationSingleUpsert apiReq identifier conf sCache
-      resultSet <- runQuery roleIsoLvl (Plan.mrTxMode mrPlan) $ Query.singleUpsertQuery mrPlan apiReq conf
-      pgrst <- liftEither $ Response.singleUpsertResponse mrPlan apiReq resultSet serverTimingParams
-      return $ pgrstResponse pgrst
+      (planTime', mrPlan) <- withTiming $ liftEither $ Plan.mutateReadPlan MutationSingleUpsert apiReq identifier conf sCache
+      (rsTime', resultSet) <- withTiming $ runQuery roleIsoLvl (Plan.mrTxMode mrPlan) $ Query.singleUpsertQuery mrPlan apiReq conf
+      (renderTime', pgrst) <- withTiming $ liftEither $ Response.singleUpsertResponse mrPlan apiReq resultSet
+      let metrics = Map.fromList [(SMPlan, planTime'), (SMQuery, rsTime'), (SMRender, renderTime'), jwtTime]
+      return $ pgrstResponse metrics pgrst

    (ActionMutate MutationDelete, TargetIdent identifier) -> do
-      mrPlan <- liftEither $ Plan.mutateReadPlan MutationDelete apiReq identifier conf sCache
-      resultSet <- runQuery roleIsoLvl (Plan.mrTxMode mrPlan) $ Query.deleteQuery mrPlan apiReq conf
-      pgrst <- liftEither $ Response.deleteResponse mrPlan apiReq resultSet serverTimingParams
-      return $ pgrstResponse pgrst
+      (planTime', mrPlan) <- withTiming $ liftEither $ Plan.mutateReadPlan MutationDelete apiReq identifier conf sCache
+      (rsTime', resultSet) <- withTiming $ runQuery roleIsoLvl (Plan.mrTxMode mrPlan) $ Query.deleteQuery mrPlan apiReq conf
+      (renderTime', pgrst) <- withTiming $ liftEither $ Response.deleteResponse mrPlan apiReq resultSet
+      let metrics = Map.fromList [(SMPlan, planTime'), (SMQuery, rsTime'), (SMRender, renderTime'), jwtTime]
+      return $ pgrstResponse metrics  pgrst

    (ActionInvoke invMethod, TargetProc identifier _) -> do
-      cPlan <- liftEither $ Plan.callReadPlan identifier conf sCache apiReq invMethod
-      resultSet <- runQuery (fromMaybe roleIsoLvl $ pdIsoLvl (Plan.crProc cPlan))(Plan.crTxMode cPlan) $ Query.invokeQuery (Plan.crProc cPlan) cPlan apiReq conf pgVer
-      pgrst <- liftEither $ Response.invokeResponse cPlan invMethod (Plan.crProc cPlan) apiReq resultSet serverTimingParams
-      return $ pgrstResponse pgrst
+      (planTime', cPlan) <- withTiming $ liftEither $ Plan.callReadPlan identifier conf sCache apiReq invMethod
+      (rsTime', resultSet) <- withTiming $ runQuery (fromMaybe roleIsoLvl $ pdIsoLvl (Plan.crProc cPlan)) (Plan.crTxMode cPlan) $ Query.invokeQuery (Plan.crProc cPlan) cPlan apiReq conf pgVer
+      (renderTime', pgrst) <- withTiming $ liftEither $ Response.invokeResponse cPlan invMethod (Plan.crProc cPlan) apiReq resultSet
+      let metrics = Map.fromList [(SMPlan, planTime'), (SMQuery, rsTime'), (SMRender, renderTime'), jwtTime]
+      return $ pgrstResponse metrics  pgrst

    (ActionInspect headersOnly, TargetDefaultSpec tSchema) -> do
-      iPlan <- liftEither $ Plan.inspectPlan conf apiReq
-      oaiResult <- runQuery roleIsoLvl (Plan.ipTxmode iPlan) $ Query.openApiQuery sCache pgVer conf tSchema
-      pgrst <- liftEither $ Response.openApiResponse (T.decodeUtf8 prettyVersion, docsVersion) headersOnly oaiResult conf sCache iSchema iNegotiatedByProfile
-      return $ pgrstResponse pgrst
+      (planTime', iPlan) <- withTiming $ liftEither $ Plan.inspectPlan conf apiReq
+      (rsTime', oaiResult) <- withTiming $ runQuery roleIsoLvl (Plan.ipTxmode iPlan) $ Query.openApiQuery sCache pgVer conf tSchema
+      (renderTime', pgrst) <- withTiming $ liftEither $ Response.openApiResponse (T.decodeUtf8 prettyVersion, docsVersion) headersOnly oaiResult conf sCache iSchema iNegotiatedByProfile
+      let metrics = Map.fromList [(SMPlan, planTime'), (SMQuery, rsTime'), (SMRender, renderTime'), jwtTime]
+      return $ pgrstResponse metrics  pgrst

    (ActionInfo, TargetIdent identifier) -> do
-      pgrst <- liftEither $ Response.infoIdentResponse identifier sCache
-      return $ pgrstResponse pgrst
+      (renderTime', pgrst) <- withTiming $ liftEither $ Response.infoIdentResponse identifier sCache
+      let metrics = Map.fromList [(SMRender, renderTime'), jwtTime]
+      return $ pgrstResponse metrics  pgrst

    (ActionInfo, TargetProc identifier _) -> do
-      cPlan <- liftEither $ Plan.callReadPlan identifier conf sCache apiReq ApiRequest.InvHead
-      pgrst <- liftEither $ Response.infoProcResponse (Plan.crProc cPlan)
-      return $ pgrstResponse pgrst
+      (planTime', cPlan) <- withTiming $ liftEither $ Plan.callReadPlan identifier conf sCache apiReq ApiRequest.InvHead
+      (renderTime', pgrst) <- withTiming $ liftEither $ Response.infoProcResponse (Plan.crProc cPlan)
+      let metrics = Map.fromList [(SMPlan, planTime'), (SMRender, renderTime'), jwtTime]
+      return $ pgrstResponse metrics  pgrst

    (ActionInfo, TargetDefaultSpec _) -> do
-      pgrst <- liftEither Response.infoRootResponse
-      return $ pgrstResponse pgrst
+      (renderTime', pgrst) <- withTiming $ liftEither Response.infoRootResponse
+      let metrics = Map.fromList [(SMRender, renderTime'), jwtTime]
+      return $ pgrstResponse metrics  pgrst

    _ ->
      -- This is unreachable as the ApiRequest.hs rejects it before
@@ -241,8 +255,16 @@ handleRequest AuthResult{..} conf appState authenticated prepared pgVer apiReq@A
        Query.runPreReq conf
        query

-    pgrstResponse :: Response.PgrstResponse -> Wai.Response
-    pgrstResponse (Response.PgrstResponse st hdrs bod) = Wai.responseLBS st hdrs bod
+    pgrstResponse :: ServerTimingData -> Response.PgrstResponse -> Wai.Response
+    pgrstResponse timings (Response.PgrstResponse st hdrs bod) = Wai.responseLBS st (hdrs ++ ([renderServerTimingHeader timings | configDbPlanEnabled conf])) bod
+
+    withTiming f = if configDbPlanEnabled conf
+        then do
+          (t, r) <- timeItT f
+          pure (Just t, r)
+        else do
+          r <- f
+          pure (Nothing, r)

 traceHeaderMiddleware :: AppConfig -> Wai.Middleware
 traceHeaderMiddleware AppConfig{configServerTraceHeader} app req respond =