risingwave_compute/rpc/service/

monitor_service.rs

// Copyright 2022 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::{BTreeMap, HashMap, HashSet};
use std::time::Duration;

use foyer::{HybridCache, TracingOptions};
use prometheus::core::Collector;
use prometheus::proto::Metric;
use risingwave_common::config::{MetricLevel, ServerConfig};
use risingwave_common_heap_profiling::ProfileServiceImpl;
use risingwave_hummock_sdk::HummockSstableObjectId;
use risingwave_jni_core::jvm_runtime::dump_jvm_stack_traces;
use risingwave_pb::monitor_service::monitor_service_server::MonitorService;
use risingwave_pb::monitor_service::stack_trace_request::ActorTracesFormat;
use risingwave_pb::monitor_service::{
    AnalyzeHeapRequest, AnalyzeHeapResponse, ChannelStats, FragmentStats, GetProfileStatsRequest,
    GetProfileStatsResponse, GetStreamingStatsRequest, GetStreamingStatsResponse,
    HeapProfilingRequest, HeapProfilingResponse, ListHeapProfilingRequest,
    ListHeapProfilingResponse, ProfilingRequest, ProfilingResponse, RelationStats,
    StackTraceRequest, StackTraceResponse, TieredCacheTracingRequest, TieredCacheTracingResponse,
};
use risingwave_storage::hummock::compactor::await_tree_key::Compaction;
use risingwave_storage::hummock::{Block, Sstable, SstableBlockIndex};
use risingwave_stream::executor::monitor::global_streaming_metrics;
use risingwave_stream::task::LocalStreamManager;
use risingwave_stream::task::await_tree_key::{Actor, BarrierAwait};
use thiserror_ext::AsReport;
use tonic::{Request, Response, Status};

type MetaCache = HybridCache<HummockSstableObjectId, Box<Sstable>>;
type BlockCache = HybridCache<SstableBlockIndex, Box<Block>>;

#[derive(Clone)]
pub struct MonitorServiceImpl {
    stream_mgr: LocalStreamManager,
    profile_service: ProfileServiceImpl,
    meta_cache: Option<MetaCache>,
    block_cache: Option<BlockCache>,
}

impl MonitorServiceImpl {
    pub fn new(
        stream_mgr: LocalStreamManager,
        server_config: ServerConfig,
        meta_cache: Option<MetaCache>,
        block_cache: Option<BlockCache>,
    ) -> Self {
        Self {
            stream_mgr,
            profile_service: ProfileServiceImpl::new(server_config),
            meta_cache,
            block_cache,
        }
    }
}

#[async_trait::async_trait]
impl MonitorService for MonitorServiceImpl {
    async fn stack_trace(
        &self,
        request: Request<StackTraceRequest>,
    ) -> Result<Response<StackTraceResponse>, Status> {
        let req = request.into_inner();

        let actor_traces = if let Some(reg) = self.stream_mgr.await_tree_reg() {
            reg.collect::<Actor>()
                .into_iter()
                .map(|(k, v)| {
                    (
                        k.0.as_raw_id(),
                        if req.actor_traces_format == ActorTracesFormat::Text as i32 {
                            v.to_string()
                        } else {
                            serde_json::to_string(&v).unwrap()
                        },
                    )
                })
                .collect()
        } else {
            Default::default()
        };

        let barrier_traces = if let Some(reg) = self.stream_mgr.await_tree_reg() {
            reg.collect::<BarrierAwait>()
                .into_iter()
                .map(|(k, v)| (k.prev_epoch, v.to_string()))
                .collect()
        } else {
            Default::default()
        };

        let rpc_traces = if let Some(reg) = self.stream_mgr.await_tree_reg() {
            reg.collect::<GrpcCall>()
                .into_iter()
                .map(|(k, v)| (k.desc, v.to_string()))
                .collect()
        } else {
            Default::default()
        };

        let compaction_task_traces = if let Some(hummock) =
            self.stream_mgr.env.state_store().as_hummock()
            && let Some(m) = hummock.compaction_await_tree_reg()
        {
            m.collect::<Compaction>()
                .into_iter()
                .map(|(k, v)| (format!("{k:?}"), v.to_string()))
                .collect()
        } else {
            Default::default()
        };

        let barrier_worker_state = self.stream_mgr.inspect_barrier_state().await?;

        let jvm_stack_traces = match dump_jvm_stack_traces() {
            Ok(None) => None,
            Err(err) => Some(err.as_report().to_string()),
            Ok(Some(stack_traces)) => Some(stack_traces),
        };

        Ok(Response::new(StackTraceResponse {
            actor_traces,
            rpc_traces,
            compaction_task_traces,
            inflight_barrier_traces: barrier_traces,
            barrier_worker_state: BTreeMap::from_iter([(
                self.stream_mgr.env.worker_id(),
                barrier_worker_state,
            )]),
            jvm_stack_traces: match jvm_stack_traces {
                Some(stack_traces) => {
                    BTreeMap::from_iter([(self.stream_mgr.env.worker_id(), stack_traces)])
                }
                None => BTreeMap::new(),
            },
            meta_traces: Default::default(),
            node_errors: Default::default(),
        }))
    }

    async fn profiling(
        &self,
        request: Request<ProfilingRequest>,
    ) -> Result<Response<ProfilingResponse>, Status> {
        self.profile_service.profiling(request).await
    }

    async fn heap_profiling(
        &self,
        request: Request<HeapProfilingRequest>,
    ) -> Result<Response<HeapProfilingResponse>, Status> {
        self.profile_service.heap_profiling(request).await
    }

    async fn list_heap_profiling(
        &self,
        _request: Request<ListHeapProfilingRequest>,
    ) -> Result<Response<ListHeapProfilingResponse>, Status> {
        self.profile_service.list_heap_profiling(_request).await
    }

    async fn analyze_heap(
        &self,
        request: Request<AnalyzeHeapRequest>,
    ) -> Result<Response<AnalyzeHeapResponse>, Status> {
        self.profile_service.analyze_heap(request).await
    }

    async fn get_profile_stats(
        &self,
        request: Request<GetProfileStatsRequest>,
    ) -> Result<Response<GetProfileStatsResponse>, Status> {
        let metrics = global_streaming_metrics(MetricLevel::Info);
        let inner = request.into_inner();
        let executor_ids = &inner.executor_ids;
        let fragment_ids = HashSet::from_iter(inner.dispatcher_fragment_ids);
        let stream_node_output_row_count = metrics
            .mem_stream_node_output_row_count
            .collect(executor_ids);
        let stream_node_output_blocking_duration_ns = metrics
            .mem_stream_node_output_blocking_duration_ns
            .collect(executor_ids);

        // Collect count metrics by fragment_ids
        fn collect_by_fragment_ids<T: Collector>(
            m: &T,
            fragment_ids: &HashSet<FragmentId>,
        ) -> HashMap<FragmentId, 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().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_ns,
            dispatch_fragment_output_row_count,
            dispatch_fragment_output_blocking_duration_ns,
        }))
    }

    async fn get_streaming_stats(
        &self,
        _request: Request<GetStreamingStatsRequest>,
    ) -> Result<Response<GetStreamingStatsResponse>, Status> {
        let metrics = global_streaming_metrics(MetricLevel::Info);

        fn collect<T: Collector>(m: &T) -> Vec<Metric> {
            m.collect().into_iter().next().unwrap().take_metric()
        }

        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_infallible(m, "fragment_id");
                let count = m.get_gauge().value() as u32;
                (fragment_id, count)
            })
            .collect();

        let mut fragment_stats: HashMap<u32, FragmentStats> = HashMap::new();
        for (&fragment_id, &actor_count) in &actor_count {
            fragment_stats.insert(
                fragment_id,
                FragmentStats {
                    actor_count,
                    current_epoch: 0,
                },
            );
        }

        let actor_current_epoch = collect(&metrics.actor_current_epoch);
        for m in &actor_current_epoch {
            let fragment_id: u32 = get_label_infallible(m, "fragment_id");
            let epoch = m.get_gauge().value() as u64;
            if let Some(s) = fragment_stats.get_mut(&fragment_id) {
                s.current_epoch = if s.current_epoch == 0 {
                    epoch
                } else {
                    u64::min(s.current_epoch, epoch)
                }
            } else {
                warn!(
                    fragment_id = fragment_id,
                    "Miss corresponding actor count metrics"
                );
            }
        }

        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_infallible(m, "table_id");
            let epoch = m.get_gauge().value() as u64;
            if let Some(s) = relation_stats.get_mut(&table_id) {
                s.current_epoch = if s.current_epoch == 0 {
                    epoch
                } else {
                    u64::min(s.current_epoch, epoch)
                };
                s.actor_count += 1;
            } else {
                relation_stats.insert(
                    table_id,
                    RelationStats {
                        actor_count: 1,
                        current_epoch: epoch,
                    },
                );
            }
        }

        let mut channel_stats: BTreeMap<String, ChannelStats> = BTreeMap::new();

        for metric in actor_output_buffer_blocking_duration_ns {
            let fragment_id: u32 = get_label_infallible(&metric, "fragment_id");
            let downstream_fragment_id: u32 =
                get_label_infallible(&metric, "downstream_fragment_id");

            let actor_count_to_add =
                if get_label_infallible::<String>(&metric, "actor_id").is_empty() {
                    match actor_count.get(&fragment_id) {
                        Some(&count) => count,
                        None => {
                            // Metrics can be momentarily inconsistent (or stale) across families.
                            // Skip this channel instead of crashing the compute node.
                            warn!(
                                fragment_id = fragment_id,
                                downstream_fragment_id = downstream_fragment_id,
                                "Miss corresponding actor count metrics"
                            );
                            continue;
                        }
                    }
                } else {
                    1
                };

            let key = format!("{}_{}", fragment_id, downstream_fragment_id);
            let channel_stat = channel_stats.entry(key).or_insert_with(|| ChannelStats {
                actor_count: 0,
                output_blocking_duration: 0.,
                recv_row_count: 0,
                send_row_count: 0,
            });

            // 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 += actor_count_to_add;
            channel_stat.output_blocking_duration += metric.get_counter().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_infallible(&metric, "fragment_id");

            // Find out and write to all downstream channels
            let key_prefix = format!("{}_", fragment_id);
            let key_range_end = format!("{}`", fragment_id); // '`' is next to `_`
            for (_, s) in channel_stats.range_mut(key_prefix..key_range_end) {
                s.send_row_count += metric.get_counter().value() as u64;
            }
        }

        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_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) {
                s.recv_row_count += metric.get_counter().value() as u64;
            }
        }

        let channel_stats = channel_stats.into_iter().collect();
        Ok(Response::new(GetStreamingStatsResponse {
            channel_stats,
            fragment_stats,
            relation_stats,
        }))
    }

    async fn tiered_cache_tracing(
        &self,
        request: Request<TieredCacheTracingRequest>,
    ) -> Result<Response<TieredCacheTracingResponse>, Status> {
        let req = request.into_inner();

        tracing::info!("Update tiered cache tracing config: {req:?}");

        if let Some(cache) = &self.meta_cache {
            if req.enable {
                cache.enable_tracing();
            } else {
                cache.disable_tracing();
            }
            let mut options = TracingOptions::new();
            if let Some(threshold) = req.record_hybrid_insert_threshold_ms {
                options = options
                    .with_record_hybrid_insert_threshold(Duration::from_millis(threshold as _));
            }
            if let Some(threshold) = req.record_hybrid_get_threshold_ms {
                options =
                    options.with_record_hybrid_get_threshold(Duration::from_millis(threshold as _));
            }
            if let Some(threshold) = req.record_hybrid_remove_threshold_ms {
                options = options
                    .with_record_hybrid_remove_threshold(Duration::from_millis(threshold as _));
            }
            if let Some(threshold) = req.record_hybrid_fetch_threshold_ms {
                options = options.with_record_hybrid_get_or_fetch_threshold(Duration::from_millis(
                    threshold as _,
                ));
            }
            cache.update_tracing_options(options);
        }

        if let Some(cache) = &self.block_cache {
            if req.enable {
                cache.enable_tracing();
            } else {
                cache.disable_tracing();
            }
            let mut options = TracingOptions::new();
            if let Some(threshold) = req.record_hybrid_insert_threshold_ms {
                options = options
                    .with_record_hybrid_insert_threshold(Duration::from_millis(threshold as _));
            }
            if let Some(threshold) = req.record_hybrid_get_threshold_ms {
                options =
                    options.with_record_hybrid_get_threshold(Duration::from_millis(threshold as _));
            }
            if let Some(threshold) = req.record_hybrid_remove_threshold_ms {
                options = options
                    .with_record_hybrid_remove_threshold(Duration::from_millis(threshold as _));
            }
            if let Some(threshold) = req.record_hybrid_fetch_threshold_ms {
                options = options.with_record_hybrid_get_or_fetch_threshold(Duration::from_millis(
                    threshold as _,
                ));
            }
            cache.update_tracing_options(options);
        }

        Ok(Response::new(TieredCacheTracingResponse::default()))
    }
}

pub use grpc_middleware::*;
use risingwave_common::metrics::get_label_infallible;
use risingwave_pb::id::FragmentId;

pub mod grpc_middleware {
    use std::sync::Arc;
    use std::sync::atomic::{AtomicU64, Ordering};
    use std::task::{Context, Poll};

    use either::Either;
    use futures::Future;
    use tonic::body::Body;
    use tower::{Layer, Service};

    /// Manages the await-trees of `gRPC` requests that are currently served by the compute node.
    pub type AwaitTreeRegistryRef = await_tree::Registry;

    /// Await-tree key type for gRPC calls.
    #[derive(Debug, Clone, PartialEq, Eq, Hash)]
    pub struct GrpcCall {
        pub desc: String,
    }

    #[derive(Clone)]
    pub struct AwaitTreeMiddlewareLayer {
        registry: Option<AwaitTreeRegistryRef>,
    }

    impl AwaitTreeMiddlewareLayer {
        pub fn new(registry: AwaitTreeRegistryRef) -> Self {
            Self {
                registry: Some(registry),
            }
        }

        pub fn new_optional(registry: Option<AwaitTreeRegistryRef>) -> Self {
            Self { registry }
        }
    }

    impl<S> Layer<S> for AwaitTreeMiddlewareLayer {
        type Service = AwaitTreeMiddleware<S>;

        fn layer(&self, service: S) -> Self::Service {
            AwaitTreeMiddleware {
                inner: service,
                registry: self.registry.clone(),
                next_id: Default::default(),
            }
        }
    }

    #[derive(Clone)]
    pub struct AwaitTreeMiddleware<S> {
        inner: S,
        registry: Option<AwaitTreeRegistryRef>,
        next_id: Arc<AtomicU64>,
    }

    impl<S> Service<http::Request<Body>> for AwaitTreeMiddleware<S>
    where
        S: Service<http::Request<Body>> + Clone,
    {
        type Error = S::Error;
        type Response = S::Response;

        type Future = impl Future<Output = Result<Self::Response, Self::Error>>;

        fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            self.inner.poll_ready(cx)
        }

        fn call(&mut self, req: http::Request<Body>) -> Self::Future {
            let Some(registry) = self.registry.clone() else {
                return Either::Left(self.inner.call(req));
            };

            // This is necessary because tonic internally uses `tower::buffer::Buffer`.
            // See https://github.com/tower-rs/tower/issues/547#issuecomment-767629149
            // for details on why this is necessary
            let clone = self.inner.clone();
            let mut inner = std::mem::replace(&mut self.inner, clone);

            let id = self.next_id.fetch_add(1, Ordering::SeqCst);
            let desc = if let Some(authority) = req.uri().authority() {
                format!("{authority} - {id}")
            } else {
                format!("?? - {id}")
            };
            let key = GrpcCall { desc };

            Either::Right(async move {
                let root = registry.register(key, req.uri().path());

                root.instrument(inner.call(req)).await
            })
        }
    }

    #[cfg(not(madsim))]
    impl<S: tonic::server::NamedService> tonic::server::NamedService for AwaitTreeMiddleware<S> {
        const NAME: &'static str = S::NAME;
    }
}