server: add dynamic labels to Prometheus metric hasura_websocket_messages_sent_bytes_total

PR-URL: https://github.com/hasura/graphql-engine-mono/pull/9920
GitOrigin-RevId: 7fb49f07d7914d08576dd0062fbea406b897e45e

docs/docs/observability/enterprise-edition/prometheus/metrics.mdx

@@ -303,6 +303,26 @@ consider investigating the subscription query and see if indexes can help improv
 | Type   | Histogram<br /><br />Buckets: 0.000001, 0.0001, 0.01, 0.1, 0.3, 1, 3, 10, 30, 100          |
 | Labels | `subscription_kind`: streaming \| live-query, `operation_name`, `parameterized_query_hash` |
 
+#### WebSocket Egress
+
+The total size of WebSocket messages sent in bytes.
+
+|        |                                              |
+| ------ | -------------------------------------------- |
+| Name   | `hasura_websocket_messages_sent_bytes_total` |
+| Type   | Counter                                      |
+| Labels | `operation_name`, `parameterized_query_hash` |
+
+#### WebSocket Ingress
+
+The total size of WebSocket messages received in bytes.
+
+|        |                                                  |
+| ------ | ------------------------------------------------ |
+| Name   | `hasura_websocket_messages_received_bytes_total` |
+| Type   | Counter                                          |
+| Labels | none                                             |
+
 #### Websocket Message Queue Time
 
 The time for which a websocket message remains queued in the GraphQL engine's websocket queue.

server/src-lib/Hasura/GraphQL/Execute/Subscription/State.hs

@@ -250,7 +250,7 @@ addLiveQuery
       liftIO $ Prometheus.Gauge.inc $ submActiveLiveQueryPollers $ pmSubscriptionMetrics $ prometheusMetrics
 
     liftIO $ EKG.Gauge.inc $ smActiveSubscriptions serverMetrics
-    let promMetricGranularLabel = SubscriptionLabel liveQuerySubscriptionLabel (Just $ DynamicSubscriptionLabel parameterizedQueryHash operationName)
+    let promMetricGranularLabel = SubscriptionLabel liveQuerySubscriptionLabel (Just $ DynamicSubscriptionLabel (Just parameterizedQueryHash) operationName)
         promMetricLabel = SubscriptionLabel liveQuerySubscriptionLabel Nothing
     let numSubscriptionMetric = submActiveSubscriptions $ pmSubscriptionMetrics $ prometheusMetrics
     recordMetricWithLabel
@@ -375,7 +375,7 @@ addStreamSubscriptionQuery
       EKG.Gauge.inc $ smActiveSubscriptions serverMetrics
       EKG.Gauge.inc $ smActiveStreamingSubscriptions serverMetrics
 
-    let promMetricGranularLabel = SubscriptionLabel streamingSubscriptionLabel (Just $ DynamicSubscriptionLabel parameterizedQueryHash operationName)
+    let promMetricGranularLabel = SubscriptionLabel streamingSubscriptionLabel (Just $ DynamicSubscriptionLabel (Just parameterizedQueryHash) operationName)
         promMetricLabel = SubscriptionLabel streamingSubscriptionLabel Nothing
         numSubscriptionMetric = submActiveSubscriptions $ pmSubscriptionMetrics $ prometheusMetrics
     recordMetricWithLabel
@@ -455,7 +455,7 @@ removeLiveQuery logger serverMetrics prometheusMetrics lqState lqId@(SubscriberD
           <*> TMap.null newOps
       when cohortIsEmpty $ TMap.delete cohortId cohortMap
       handlerIsEmpty <- TMap.null cohortMap
-      let promMetricGranularLabel = SubscriptionLabel liveQuerySubscriptionLabel (Just $ DynamicSubscriptionLabel parameterizedQueryHash maybeOperationName)
+      let promMetricGranularLabel = SubscriptionLabel liveQuerySubscriptionLabel (Just $ DynamicSubscriptionLabel (Just parameterizedQueryHash) maybeOperationName)
           promMetricLabel = SubscriptionLabel liveQuerySubscriptionLabel Nothing
       -- when there is no need for handler i.e, this happens to be the last
       -- operation, take the ref for the polling thread to cancel it
@@ -554,7 +554,7 @@ removeStreamingQuery logger serverMetrics prometheusMetrics subscriptionState (S
           <*> TMap.null newOps
       when cohortIsEmpty $ TMap.delete currentCohortId cohortMap
       handlerIsEmpty <- TMap.null cohortMap
-      let promMetricGranularLabel = SubscriptionLabel streamingSubscriptionLabel (Just $ DynamicSubscriptionLabel parameterizedQueryHash maybeOperationName)
+      let promMetricGranularLabel = SubscriptionLabel streamingSubscriptionLabel (Just $ DynamicSubscriptionLabel (Just parameterizedQueryHash) maybeOperationName)
           promMetricLabel = SubscriptionLabel streamingSubscriptionLabel Nothing
       -- when there is no need for handler i.e,
       -- operation, take the ref for the polling thread to cancel it

server/src-lib/Hasura/GraphQL/Transport/WebSocket/Server.hs

@@ -72,9 +72,11 @@ import Hasura.Server.Auth (AuthMode, compareAuthMode)
 import Hasura.Server.Cors (CorsPolicy)
 import Hasura.Server.Init.Config (AllowListStatus (..), WSConnectionInitTimeout (..))
 import Hasura.Server.Prometheus
-  ( PrometheusMetrics (..),
+  ( DynamicSubscriptionLabel (..),
+    PrometheusMetrics (..),
+    recordMetricWithLabel,
   )
-import Hasura.Server.Types (ExperimentalFeature (..))
+import Hasura.Server.Types (ExperimentalFeature (..), MonadGetPolicies (runGetPrometheusMetricsGranularity))
 import ListT qualified
 import Network.Wai.Extended (IpAddress)
 import Network.Wai.Handler.Warp qualified as Warp
@@ -83,6 +85,7 @@ import Refined (unrefine)
 import StmContainers.Map qualified as STMMap
 import System.IO.Error qualified as E
 import System.Metrics.Prometheus.Counter qualified as Prometheus.Counter
+import System.Metrics.Prometheus.CounterVector qualified as CounterVector
 import System.Metrics.Prometheus.Histogram qualified as Prometheus.Histogram
 import System.TimeManager qualified as TM
 
@@ -497,7 +500,7 @@ websocketConnectionReaper getLatestConfig getSchemaCache ws@(WSServer _ userConf
           | otherwise -> pure ()
 
 createServerApp ::
-  (MonadIO m, MC.MonadBaseControl IO m, LA.Forall (LA.Pure m), MonadWSLog m) =>
+  (MonadIO m, MC.MonadBaseControl IO m, LA.Forall (LA.Pure m), MonadWSLog m, MonadGetPolicies m) =>
   IO MetricsConfig ->
   WSConnectionInitTimeout ->
   WSServer a ->
@@ -630,10 +633,18 @@ createServerApp getMetricsConfig wsConnInitTimeout (WSServer logger@(L.Logger wr
                   (messageWriteTime, _) <- liftIO $ withElapsedTime $ WS.sendTextData conn msg
                   let messageLength = BL.length msg
                       messageDetails = MessageDetails (SB.fromLBS msg) messageLength
+                      parameterizedQueryHash = wsInfo >>= _wseiParameterizedQueryHash
+                      operationName = wsInfo >>= _wseiOperationName
+                      promMetricGranularLabel = DynamicSubscriptionLabel parameterizedQueryHash operationName
+                      promMetricLabel = DynamicSubscriptionLabel Nothing Nothing
+                      websocketBytesSentMetric = pmWebSocketBytesSent prometheusMetrics
+                  granularPrometheusMetricsState <- runGetPrometheusMetricsGranularity
                   liftIO $ do
-                    Prometheus.Counter.add
-                      (pmWebSocketBytesSent prometheusMetrics)
-                      messageLength
+                    recordMetricWithLabel
+                      granularPrometheusMetricsState
+                      True
+                      (CounterVector.add websocketBytesSentMetric promMetricGranularLabel messageLength)
+                      (CounterVector.add websocketBytesSentMetric promMetricLabel messageLength)
                     Prometheus.Histogram.observe
                       (pmWebsocketMsgQueueTimeSeconds prometheusMetrics)
                       messageQueueTime

server/src-lib/Hasura/Server/Prometheus.hs

@@ -71,7 +71,7 @@ data PrometheusMetrics = PrometheusMetrics
     pmGraphQLRequestMetrics :: GraphQLRequestMetrics,
     pmEventTriggerMetrics :: EventTriggerMetrics,
     pmWebSocketBytesReceived :: Counter,
-    pmWebSocketBytesSent :: Counter,
+    pmWebSocketBytesSent :: CounterVector DynamicSubscriptionLabel,
     pmActionBytesReceived :: Counter,
     pmActionBytesSent :: Counter,
     pmScheduledTriggerMetrics :: ScheduledTriggerMetrics,
@@ -157,7 +157,7 @@ makeDummyPrometheusMetrics = do
   pmGraphQLRequestMetrics <- makeDummyGraphQLRequestMetrics
   pmEventTriggerMetrics <- makeDummyEventTriggerMetrics
   pmWebSocketBytesReceived <- Counter.new
-  pmWebSocketBytesSent <- Counter.new
+  pmWebSocketBytesSent <- CounterVector.new
   pmActionBytesReceived <- Counter.new
   pmActionBytesSent <- Counter.new
   pmScheduledTriggerMetrics <- makeDummyScheduledTriggerMetrics
@@ -336,7 +336,7 @@ liveQuerySubscriptionLabel :: SubscriptionKindLabel
 liveQuerySubscriptionLabel = SubscriptionKindLabel "live-query"
 
 data DynamicSubscriptionLabel = DynamicSubscriptionLabel
-  { _dslParamQueryHash :: ParameterizedQueryHash,
+  { _dslParamQueryHash :: Maybe ParameterizedQueryHash,
     _dslOperationName :: Maybe OperationName
   }
   deriving stock (Generic, Ord, Eq)
@@ -344,7 +344,7 @@ data DynamicSubscriptionLabel = DynamicSubscriptionLabel
 instance ToLabels DynamicSubscriptionLabel where
   toLabels (DynamicSubscriptionLabel hash opName) =
     Map.fromList
-      $ [("parameterized_query_hash", bsToTxt $ unParamQueryHash hash)]
+      $ maybe [] (\pqh -> [("parameterized_query_hash", bsToTxt $ unParamQueryHash pqh)]) hash
       <> maybe [] (\op -> [("operation_name", G.unName $ _unOperationName op)]) opName
 
 data SubscriptionLabel = SubscriptionLabel
@@ -416,7 +416,7 @@ recordSubcriptionMetric getMetricState alwaysObserve operationNamesMap parameter
   -- if no operation names are present, then emit metric with only param query hash as dynamic label
   if (null operationNamesMap)
     then do
-      let promMetricGranularLabel = SubscriptionLabel subscriptionKind (Just $ DynamicSubscriptionLabel parameterizedQueryHash Nothing)
+      let promMetricGranularLabel = SubscriptionLabel subscriptionKind (Just $ DynamicSubscriptionLabel (Just parameterizedQueryHash) Nothing)
           promMetricLabel = SubscriptionLabel subscriptionKind Nothing
       recordMetricWithLabel
         getMetricState
@@ -427,7 +427,7 @@ recordSubcriptionMetric getMetricState alwaysObserve operationNamesMap parameter
     do
       let operationNames = HashMap.keys operationNamesMap
       for_ operationNames $ \opName -> do
-        let promMetricGranularLabel = SubscriptionLabel subscriptionKind (Just $ DynamicSubscriptionLabel parameterizedQueryHash opName)
+        let promMetricGranularLabel = SubscriptionLabel subscriptionKind (Just $ DynamicSubscriptionLabel (Just parameterizedQueryHash) opName)
             promMetricLabel = SubscriptionLabel subscriptionKind Nothing
         recordMetricWithLabel
           getMetricState