common_runtime/
metrics.rs

1// Copyright 2023 Greptime Team
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
15//! Runtime metrics
16use std::collections::BTreeMap;
17use std::sync::Mutex;
18use std::time::Duration;
19
20use catio::Scheduler;
21use lazy_static::lazy_static;
22use prometheus::core::{Collector, Desc};
23use prometheus::proto::MetricFamily;
24use prometheus::*;
25
26use crate::global::{QUERY_TASK_CLASS, WRITE_TASK_CLASS};
27
28pub const THREAD_NAME_LABEL: &str = "thread_name";
29
30lazy_static! {
31    pub static ref METRIC_RUNTIME_THREADS_ALIVE: IntGaugeVec = register_int_gauge_vec!(
32        "greptime_runtime_threads_alive",
33        "runtime threads alive",
34        &[THREAD_NAME_LABEL]
35    )
36    .unwrap();
37    pub static ref METRIC_RUNTIME_THREADS_IDLE: IntGaugeVec = register_int_gauge_vec!(
38        "greptime_runtime_threads_idle",
39        "runtime threads idle",
40        &[THREAD_NAME_LABEL]
41    )
42    .unwrap();
43}
44
45#[derive(Clone, Default)]
46struct ClassSnapshot {
47    polls: u64,
48    total_admission_wait: Duration,
49}
50
51struct WorkloadSchedulerCollector {
52    scheduler: Scheduler,
53    enabled: IntGauge,
54    active: IntGauge,
55    weight: IntGaugeVec,
56    queued: IntGaugeVec,
57    polls: IntCounterVec,
58    total_admission_wait: CounterVec,
59    snapshots: Mutex<BTreeMap<&'static str, ClassSnapshot>>,
60}
61
62impl WorkloadSchedulerCollector {
63    fn new(scheduler: Scheduler) -> Self {
64        let workload_label = &["workload"];
65        Self {
66            scheduler,
67            enabled: IntGauge::new(
68                "greptime_workload_scheduler_enabled",
69                "Whether the workload scheduler is enabled",
70            )
71            .unwrap(),
72            active: IntGauge::new(
73                "greptime_workload_scheduler_active_polls",
74                "Task polls admitted to Tokio but not yet completed",
75            )
76            .unwrap(),
77            weight: IntGaugeVec::new(
78                Opts::new(
79                    "greptime_workload_scheduler_weight",
80                    "Configured workload scheduler weight",
81                ),
82                workload_label,
83            )
84            .unwrap(),
85            queued: IntGaugeVec::new(
86                Opts::new(
87                    "greptime_workload_scheduler_queued_tasks",
88                    "Tasks queued in the workload scheduler",
89                ),
90                workload_label,
91            )
92            .unwrap(),
93            polls: IntCounterVec::new(
94                Opts::new(
95                    "greptime_workload_scheduler_polls_total",
96                    "Cumulative task polls admitted by the workload scheduler",
97                ),
98                workload_label,
99            )
100            .unwrap(),
101            total_admission_wait: CounterVec::new(
102                Opts::new(
103                    "greptime_workload_scheduler_admission_wait_seconds_total",
104                    "Cumulative workload scheduler admission wait time in seconds",
105                ),
106                workload_label,
107            )
108            .unwrap(),
109            snapshots: Mutex::new(BTreeMap::new()),
110        }
111    }
112
113    fn update_locked(&self, snapshots: &mut BTreeMap<&'static str, ClassSnapshot>) {
114        let stats = self.scheduler.stats();
115        self.enabled.set(i64::from(self.scheduler.is_enabled()));
116        self.active.set(
117            stats
118                .active_polls
119                .min(i64::MAX as usize)
120                .try_into()
121                .unwrap_or(i64::MAX),
122        );
123
124        for (class, workload) in [(QUERY_TASK_CLASS, "query"), (WRITE_TASK_CLASS, "write")] {
125            let class_stats = stats.classes.get(&class).cloned().unwrap_or_default();
126            let labels = &[workload];
127            self.weight
128                .with_label_values(labels)
129                .set(i64::from(class_stats.weight));
130            self.queued.with_label_values(labels).set(
131                class_stats
132                    .queued
133                    .min(i64::MAX as usize)
134                    .try_into()
135                    .unwrap_or(i64::MAX),
136            );
137            let previous = snapshots.entry(workload).or_default();
138            self.polls
139                .with_label_values(labels)
140                .inc_by(class_stats.polls.saturating_sub(previous.polls));
141            self.total_admission_wait.with_label_values(labels).inc_by(
142                class_stats
143                    .total_admission_wait
144                    .saturating_sub(previous.total_admission_wait)
145                    .as_secs_f64(),
146            );
147            *previous = ClassSnapshot {
148                polls: class_stats.polls,
149                total_admission_wait: class_stats.total_admission_wait,
150            };
151        }
152    }
153}
154
155impl Collector for WorkloadSchedulerCollector {
156    fn desc(&self) -> Vec<&Desc> {
157        let mut desc = self.enabled.desc();
158        desc.extend(self.active.desc());
159        desc.extend(self.weight.desc());
160        desc.extend(self.queued.desc());
161        desc.extend(self.polls.desc());
162        desc.extend(self.total_admission_wait.desc());
163        desc
164    }
165
166    fn collect(&self) -> Vec<MetricFamily> {
167        let mut snapshots = self.snapshots.lock().unwrap();
168        self.update_locked(&mut snapshots);
169        let mut families = self.enabled.collect();
170        families.extend(self.active.collect());
171        families.extend(self.weight.collect());
172        families.extend(self.queued.collect());
173        families.extend(self.polls.collect());
174        families.extend(self.total_admission_wait.collect());
175        families
176    }
177}
178
179pub(crate) fn register_workload_scheduler_metrics(scheduler: Scheduler) {
180    register(Box::new(WorkloadSchedulerCollector::new(scheduler)))
181        .expect("workload scheduler metrics collector registration must succeed");
182}
183
184#[cfg(test)]
185mod tests {
186    use std::collections::{BTreeMap, BTreeSet};
187    use std::time::Duration;
188
189    use prometheus::proto::{MetricFamily, MetricType};
190
191    use super::*;
192
193    fn family<'a>(families: &'a [MetricFamily], name: &str) -> &'a MetricFamily {
194        families
195            .iter()
196            .find(|family| family.name() == name)
197            .unwrap_or_else(|| panic!("missing metric family {name}"))
198    }
199
200    fn counter_values(families: &[MetricFamily]) -> BTreeMap<String, BTreeMap<String, f64>> {
201        [
202            "greptime_workload_scheduler_polls_total",
203            "greptime_workload_scheduler_admission_wait_seconds_total",
204        ]
205        .into_iter()
206        .map(|name| {
207            let values = family(families, name)
208                .get_metric()
209                .iter()
210                .map(|metric| {
211                    (
212                        metric.get_label()[0].value().to_string(),
213                        metric.get_counter().value(),
214                    )
215                })
216                .collect();
217            (name.to_string(), values)
218        })
219        .collect()
220    }
221
222    #[test]
223    fn workload_scheduler_collector_reports_class_metrics_and_deltas() {
224        let scheduler = Scheduler::builder()
225            .max_concurrent_polls(1)
226            .weight(QUERY_TASK_CLASS, 2)
227            .weight(WRITE_TASK_CLASS, 3)
228            .build();
229        scheduler.set_enabled(true);
230        let collector = WorkloadSchedulerCollector::new(scheduler.clone());
231
232        let first = collector.collect();
233        let expected = [
234            (
235                "greptime_workload_scheduler_enabled",
236                MetricType::GAUGE,
237                false,
238            ),
239            (
240                "greptime_workload_scheduler_active_polls",
241                MetricType::GAUGE,
242                false,
243            ),
244            (
245                "greptime_workload_scheduler_weight",
246                MetricType::GAUGE,
247                true,
248            ),
249            (
250                "greptime_workload_scheduler_queued_tasks",
251                MetricType::GAUGE,
252                true,
253            ),
254            (
255                "greptime_workload_scheduler_polls_total",
256                MetricType::COUNTER,
257                true,
258            ),
259            (
260                "greptime_workload_scheduler_admission_wait_seconds_total",
261                MetricType::COUNTER,
262                true,
263            ),
264        ];
265        let expected_names: BTreeSet<_> = expected.iter().map(|(name, _, _)| *name).collect();
266        let actual_names: BTreeSet<_> = first.iter().map(MetricFamily::name).collect();
267        assert_eq!(expected_names, actual_names);
268        for (name, metric_type, has_workload_label) in expected {
269            let metric_family = family(&first, name);
270            assert_eq!(metric_type, metric_family.get_field_type(), "{name}");
271            let workloads: BTreeSet<_> = metric_family
272                .get_metric()
273                .iter()
274                .flat_map(|metric| {
275                    assert_eq!(has_workload_label, !metric.get_label().is_empty());
276                    metric
277                        .get_label()
278                        .iter()
279                        .map(|label| {
280                            assert_eq!("workload", label.name());
281                            label.value()
282                        })
283                        .collect::<Vec<_>>()
284                })
285                .collect();
286            if has_workload_label {
287                assert_eq!(BTreeSet::from(["query", "write"]), workloads);
288            } else {
289                assert!(workloads.is_empty());
290            }
291        }
292
293        let second = collector.collect();
294        assert_eq!(counter_values(&first), counter_values(&second));
295
296        scheduler.set_enabled(true);
297        let runtime = tokio::runtime::Builder::new_current_thread()
298            .enable_time()
299            .build()
300            .unwrap();
301        runtime.block_on(async {
302            let query = scheduler.spawn_in(QUERY_TASK_CLASS, async {
303                tokio::time::sleep(Duration::from_millis(1)).await;
304            });
305            let write = scheduler.spawn_in(WRITE_TASK_CLASS, async {
306                tokio::time::sleep(Duration::from_millis(1)).await;
307            });
308            tokio::time::timeout(Duration::from_secs(1), async {
309                query.await.unwrap();
310                write.await.unwrap();
311            })
312            .await
313            .expect("scheduled test tasks did not complete");
314        });
315
316        let third = collector.collect();
317        let before = counter_values(&second);
318        let after = counter_values(&third);
319        for workload in ["query", "write"] {
320            let metric = "greptime_workload_scheduler_polls_total";
321            assert!(
322                after[metric][workload] > before[metric][workload],
323                "{metric} did not increase for {workload}"
324            );
325        }
326    }
327}
common_runtime/metrics.rs

common_runtime/
metrics.rs
