risingwave_ctl/cmd_impl/
profile.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashSet;
16use std::path::PathBuf;
17
18use chrono::prelude::Local;
19use clap::ValueEnum;
20use futures::future::try_join_all;
21use risingwave_pb::common::WorkerType;
22use risingwave_pb::monitor_service::ProfilingResponse;
23use risingwave_rpc_client::MonitorClientPool;
24use thiserror_ext::AsReport;
25use tokio::fs::{File, create_dir_all};
26use tokio::io::AsyncWriteExt;
27
28use crate::CtlContext;
29
30#[derive(Clone, Copy, Debug, ValueEnum)]
31#[clap(rename_all = "kebab-case")]
32pub enum ProfileWorkerType {
33    Frontend,
34    ComputeNode,
35    Compactor,
36    Meta,
37}
38
39pub async fn cpu_profile(
40    context: &CtlContext,
41    sleep_s: u64,
42    worker_types: Vec<ProfileWorkerType>,
43) -> anyhow::Result<()> {
44    let meta_client = context.meta_client().await?;
45
46    let workers = meta_client.list_worker_nodes(None).await?;
47    let target_types = selected_worker_types(&worker_types);
48    let target_nodes = workers
49        .into_iter()
50        .filter(|w| target_types.contains(&w.r#type()));
51
52    let clients = MonitorClientPool::adhoc();
53
54    let profile_root_path = std::env::var("PREFIX_PROFILING").unwrap_or_else(|_| {
55        tracing::info!("PREFIX_PROFILING is not set, using current directory");
56        "./".to_owned()
57    });
58    let profile_root_path = PathBuf::from(&profile_root_path);
59    let dir_name = Local::now().format("%Y-%m-%d-%H-%M-%S").to_string();
60    let dir_path = profile_root_path.join(dir_name);
61    create_dir_all(&dir_path).await?;
62
63    let mut profile_futs = vec![];
64
65    // FIXME: the node may not be accessible directly from risectl, we may let the meta
66    // service collect the reports from all nodes in the future.
67    for cn in target_nodes {
68        let client = clients.get(&cn).await?;
69
70        let dir_path_ref = &dir_path;
71
72        let fut = async move {
73            let response = client.profile(sleep_s).await;
74            let host_addr = cn.get_host().expect("Should have host address");
75            let node_name = format!(
76                "{}-{}-{}",
77                worker_type_label(cn.r#type()),
78                host_addr.get_host().replace('.', "-"),
79                host_addr.get_port()
80            );
81            let svg_file_name = format!("{}.svg", node_name);
82            match response {
83                Ok(ProfilingResponse { result }) => {
84                    let mut file = File::create(dir_path_ref.join(svg_file_name)).await?;
85                    file.write_all(&result).await?;
86                    file.flush().await?;
87                }
88                Err(err) => {
89                    tracing::error!(
90                        error = %err.as_report(),
91                        node_name,
92                        "Failed to get profiling result",
93                    );
94                }
95            }
96            Ok::<_, anyhow::Error>(())
97        };
98        profile_futs.push(fut);
99    }
100
101    try_join_all(profile_futs).await?;
102
103    println!("Profiling results are saved at {}", dir_path.display());
104
105    Ok(())
106}
107
108pub async fn heap_profile(
109    context: &CtlContext,
110    dir: Option<String>,
111    worker_types: Vec<ProfileWorkerType>,
112) -> anyhow::Result<()> {
113    let dir = dir.unwrap_or_default();
114    let meta_client = context.meta_client().await?;
115
116    let workers = meta_client.list_worker_nodes(None).await?;
117    let target_types = selected_worker_types(&worker_types);
118    let target_nodes = workers
119        .into_iter()
120        .filter(|w| target_types.contains(&w.r#type()));
121
122    let clients = MonitorClientPool::adhoc();
123
124    let mut profile_futs = vec![];
125
126    // FIXME: the node may not be accessible directly from risectl, we may let the meta
127    // service collect the reports from all nodes in the future.
128    for cn in target_nodes {
129        let client = clients.get(&cn).await?;
130        let dir = &dir;
131
132        let fut = async move {
133            let response = client.heap_profile(dir.clone()).await;
134            let host_addr = cn.get_host().expect("Should have host address");
135
136            let node_name = format!(
137                "{}-{}-{}",
138                worker_type_label(cn.r#type()),
139                host_addr.get_host().replace('.', "-"),
140                host_addr.get_port()
141            );
142
143            if let Err(err) = response {
144                tracing::error!(
145                    error = %err.as_report(),
146                    node_name,
147                    "Failed to dump profile",
148                );
149            }
150            Ok::<_, anyhow::Error>(())
151        };
152        profile_futs.push(fut);
153    }
154
155    try_join_all(profile_futs).await?;
156
157    println!(
158        "Profiling results are saved at {} on each target node",
159        PathBuf::from(dir).display()
160    );
161
162    Ok(())
163}
164
165fn selected_worker_types(values: &[ProfileWorkerType]) -> HashSet<WorkerType> {
166    // Use default set when no filter is provided.
167    if values.is_empty() {
168        return HashSet::from([
169            WorkerType::Frontend,
170            WorkerType::ComputeNode,
171            WorkerType::Compactor,
172            WorkerType::Meta,
173        ]);
174    }
175    values.iter().map(|value| to_worker_type(*value)).collect()
176}
177
178fn to_worker_type(value: ProfileWorkerType) -> WorkerType {
179    match value {
180        ProfileWorkerType::Frontend => WorkerType::Frontend,
181        ProfileWorkerType::ComputeNode => WorkerType::ComputeNode,
182        ProfileWorkerType::Compactor => WorkerType::Compactor,
183        ProfileWorkerType::Meta => WorkerType::Meta,
184    }
185}
186
187fn worker_type_label(worker_type: WorkerType) -> &'static str {
188    match worker_type {
189        WorkerType::Frontend => "frontend",
190        WorkerType::ComputeNode => "compute-node",
191        WorkerType::Compactor => "compactor",
192        WorkerType::Meta => "meta",
193        _ => "unknown",
194    }
195}
risingwave_ctl/cmd_impl/profile.rs

risingwave_ctl/cmd_impl/
profile.rs
