feat(frontend): support dispatcher metrics for explain analyze (#21227)

Author: kwannoel

proto/monitor_service.proto

@@ -94,12 +94,18 @@ message TieredCacheTracingRequest {
 message TieredCacheTracingResponse {}
 
 message GetProfileStatsRequest {
+  // Executors to fetch statistics for.
   repeated uint64 executor_ids = 1;
+  // Dispatchers do not have executors.
+  // We have to fetch their statistics separately.
+  repeated uint32 dispatcher_fragment_ids = 2;
 }
 
 message GetProfileStatsResponse {
   map<uint64, uint64> stream_node_output_row_count = 1;
   map<uint64, uint64> stream_node_output_blocking_duration_ms = 2;
+  map<uint32, uint64> dispatch_fragment_output_row_count = 3;
+  map<uint32, uint64> dispatch_fragment_output_blocking_duration_ns = 4;
 }
 
 service MonitorService {

src/common/metrics/src/lib.rs

@@ -19,6 +19,7 @@ use std::sync::LazyLock;
 
 use hytra::TrAdder;
 use prometheus::core::{Atomic, AtomicU64, GenericCounter, GenericGauge};
+use prometheus::proto::Metric;
 use prometheus::register_int_counter_with_registry;
 use tracing::Subscriber;
 use tracing_subscriber::Layer;
@@ -142,3 +143,16 @@ impl PartialOrd for MetricLevel {
         (*self as u8).partial_cmp(&(*other as u8))
     }
 }
+
+pub fn get_label<T: std::str::FromStr>(metric: &Metric, label: &str) -> Option<T> {
+    metric
+        .get_label()
+        .iter()
+        .find(|lp| lp.get_name() == label)
+        .and_then(|lp| lp.get_value().parse::<T>().ok())
+}
+
+// Must ensure the label exists and can be parsed into `T`
+pub fn get_label_infallible<T: std::str::FromStr>(metric: &Metric, label: &str) -> T {
+    get_label(metric, label).expect("label not found or can't be parsed")
+}

src/compute/src/rpc/service/monitor_service.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::{BTreeMap, HashMap};
+use std::collections::{BTreeMap, HashMap, HashSet};
 use std::ffi::CString;
 use std::fs;
 use std::path::Path;
@@ -293,16 +293,45 @@ impl MonitorService for MonitorServiceImpl {
         request: Request<GetProfileStatsRequest>,
     ) -> Result<Response<GetProfileStatsResponse>, Status> {
         let metrics = global_streaming_metrics(MetricLevel::Info);
-        let executor_ids = &request.into_inner().executor_ids;
+        let inner = request.into_inner();
+        let executor_ids = &inner.executor_ids;
+        let fragment_ids = HashSet::from_iter(inner.dispatcher_fragment_ids.into_iter());
         let stream_node_output_row_count = metrics
             .mem_stream_node_output_row_count
             .collect(executor_ids);
         let stream_node_output_blocking_duration_ms = metrics
             .mem_stream_node_output_blocking_duration_ms
             .collect(executor_ids);
+
+        // Collect count metrics by fragment_ids
+        fn collect_by_fragment_ids<T: Collector>(
+            m: &T,
+            fragment_ids: &HashSet<u32>,
+        ) -> HashMap<u32, u64> {
+            let mut metrics = HashMap::new();
+            for mut metric_family in m.collect() {
+                for metric in metric_family.take_metric() {
+                    let fragment_id = get_label_infallible(&metric, "fragment_id");
+                    if fragment_ids.contains(&fragment_id) {
+                        let entry = metrics.entry(fragment_id).or_insert(0);
+                        *entry += metric.get_counter().get_value() as u64;
+                    }
+                }
+            }
+            metrics
+        }
+
+        let dispatch_fragment_output_row_count =
+            collect_by_fragment_ids(&metrics.actor_out_record_cnt, &fragment_ids);
+        let dispatch_fragment_output_blocking_duration_ns = collect_by_fragment_ids(
+            &metrics.actor_output_buffer_blocking_duration_ns,
+            &fragment_ids,
+        );
         Ok(Response::new(GetProfileStatsResponse {
             stream_node_output_row_count,
             stream_node_output_blocking_duration_ms,
+            dispatch_fragment_output_row_count,
+            dispatch_fragment_output_blocking_duration_ns,
         }))
     }
 
@@ -322,27 +351,14 @@ impl MonitorService for MonitorServiceImpl {
                 .into_vec()
         }
 
-        // Must ensure the label exists and can be parsed into `T`
-        fn get_label<T: std::str::FromStr>(metric: &Metric, label: &str) -> T {
-            metric
-                .get_label()
-                .iter()
-                .find(|lp| lp.get_name() == label)
-                .unwrap()
-                .get_value()
-                .parse::<T>()
-                .ok()
-                .unwrap()
-        }
-
         let actor_output_buffer_blocking_duration_ns =
             collect(&metrics.actor_output_buffer_blocking_duration_ns);
         let actor_count = collect(&metrics.actor_count);
 
         let actor_count: HashMap<_, _> = actor_count
             .iter()
             .map(|m| {
-                let fragment_id: u32 = get_label(m, "fragment_id");
+                let fragment_id: u32 = get_label_infallible(m, "fragment_id");
                 let count = m.get_gauge().get_value() as u32;
                 (fragment_id, count)
             })
@@ -361,7 +377,7 @@ impl MonitorService for MonitorServiceImpl {
 
         let actor_current_epoch = collect(&metrics.actor_current_epoch);
         for m in &actor_current_epoch {
-            let fragment_id: u32 = get_label(m, "fragment_id");
+            let fragment_id: u32 = get_label_infallible(m, "fragment_id");
             let epoch = m.get_gauge().get_value() as u64;
             if let Some(s) = fragment_stats.get_mut(&fragment_id) {
                 s.current_epoch = if s.current_epoch == 0 {
@@ -380,7 +396,7 @@ impl MonitorService for MonitorServiceImpl {
         let mut relation_stats: HashMap<u32, RelationStats> = HashMap::new();
         let mview_current_epoch = collect(&metrics.materialize_current_epoch);
         for m in &mview_current_epoch {
-            let table_id: u32 = get_label(m, "table_id");
+            let table_id: u32 = get_label_infallible(m, "table_id");
             let epoch = m.get_gauge().get_value() as u64;
             if let Some(s) = relation_stats.get_mut(&table_id) {
                 s.current_epoch = if s.current_epoch == 0 {
@@ -403,8 +419,9 @@ impl MonitorService for MonitorServiceImpl {
         let mut channel_stats: BTreeMap<String, ChannelStats> = BTreeMap::new();
 
         for metric in actor_output_buffer_blocking_duration_ns {
-            let fragment_id: u32 = get_label(&metric, "fragment_id");
-            let downstream_fragment_id: u32 = get_label(&metric, "downstream_fragment_id");
+            let fragment_id: u32 = get_label_infallible(&metric, "fragment_id");
+            let downstream_fragment_id: u32 =
+                get_label_infallible(&metric, "downstream_fragment_id");
 
             let key = format!("{}_{}", fragment_id, downstream_fragment_id);
             let channel_stat = channel_stats.entry(key).or_insert_with(|| ChannelStats {
@@ -416,17 +433,18 @@ impl MonitorService for MonitorServiceImpl {
 
             // When metrics level is Debug, `actor_id` will be removed to reduce metrics.
             // See `src/common/metrics/src/relabeled_metric.rs`
-            channel_stat.actor_count += if get_label::<String>(&metric, "actor_id").is_empty() {
-                actor_count[&fragment_id]
-            } else {
-                1
-            };
+            channel_stat.actor_count +=
+                if get_label_infallible::<String>(&metric, "actor_id").is_empty() {
+                    actor_count[&fragment_id]
+                } else {
+                    1
+                };
             channel_stat.output_blocking_duration += metric.get_counter().get_value();
         }
 
         let actor_output_row_count = collect(&metrics.actor_out_record_cnt);
         for metric in actor_output_row_count {
-            let fragment_id: u32 = get_label(&metric, "fragment_id");
+            let fragment_id: u32 = get_label_infallible(&metric, "fragment_id");
 
             // Find out and write to all downstream channels
             let key_prefix = format!("{}_", fragment_id);
@@ -438,8 +456,8 @@ impl MonitorService for MonitorServiceImpl {
 
         let actor_input_row_count = collect(&metrics.actor_in_record_cnt);
         for metric in actor_input_row_count {
-            let upstream_fragment_id: u32 = get_label(&metric, "upstream_fragment_id");
-            let fragment_id: u32 = get_label(&metric, "fragment_id");
+            let upstream_fragment_id: u32 = get_label_infallible(&metric, "upstream_fragment_id");
+            let fragment_id: u32 = get_label_infallible(&metric, "fragment_id");
 
             let key = format!("{}_{}", upstream_fragment_id, fragment_id);
             if let Some(s) = channel_stats.get_mut(&key) {
@@ -529,6 +547,7 @@ impl MonitorService for MonitorServiceImpl {
 }
 
 pub use grpc_middleware::*;
+use risingwave_common::metrics::get_label_infallible;
 
 pub mod grpc_middleware {
     use std::sync::Arc;