mackerel_plugin/
plugin.rs

1use auto_enums::auto_enum;
2use serde_derive::{Deserialize, Serialize};
3use serde_json::json;
4use std::collections::HashMap;
5use std::io::Write;
6
7use crate::graph::Graph;
8use crate::metric::Metric;
9
10#[derive(Default, Serialize, Deserialize)]
11struct MetricValues {
12    timestamp: i64,
13    values: HashMap<String, f64>,
14}
15
16impl MetricValues {
17    fn new(timestamp: i64, values: HashMap<String, f64>) -> MetricValues {
18        MetricValues { timestamp, values }
19    }
20}
21
22/// A trait which represents a Plugin.
23///
24/// You can create a plugin by implementing `fetch_metrics` and `graph_definition`.
25pub trait Plugin {
26    fn fetch_metrics(&self) -> Result<HashMap<String, f64>, String>;
27
28    fn graph_definition(&self) -> Vec<Graph>;
29
30    fn metric_key_prefix(&self) -> String {
31        "".to_owned()
32    }
33
34    #[doc(hidden)]
35    fn output_values(&self, out: &mut dyn std::io::Write) -> Result<(), String> {
36        let now = std::time::SystemTime::now()
37            .duration_since(std::time::UNIX_EPOCH)
38            .map_err(|e| e.to_string())?;
39        let metric_values = MetricValues::new(now.as_secs() as i64, self.fetch_metrics()?);
40        let prefix = self.metric_key_prefix();
41        let graphs = self.graph_definition();
42        let has_diff = graphs.iter().any(|graph| graph.has_diff());
43        let path = self.tempfile_path(&prefix)?;
44        let prev_metric_values = if has_diff {
45            load_values(&path).unwrap_or_default()
46        } else {
47            MetricValues::default()
48        };
49        for graph in graphs {
50            for metric in graph.metrics {
51                format_values(
52                    out,
53                    &prefix,
54                    &graph.name,
55                    metric,
56                    &metric_values,
57                    &prev_metric_values,
58                );
59            }
60        }
61        if has_diff {
62            save_values(&path, &metric_values)?;
63        }
64        Ok(())
65    }
66
67    #[doc(hidden)]
68    fn tempfile_path(&self, prefix: &str) -> Result<String, String> {
69        let name = if prefix.is_empty() {
70            let arg0 = std::env::args().next().ok_or("unknown executable path")?;
71            let exec_name = std::path::Path::new(&arg0)
72                .file_name()
73                .and_then(std::ffi::OsStr::to_str)
74                .ok_or("invalid executable name")?;
75            if exec_name.starts_with("mackerel-plugin-") {
76                exec_name.to_owned()
77            } else {
78                "mackerel-plugin-".to_owned() + exec_name
79            }
80        } else {
81            "mackerel-plugin-".to_owned() + prefix
82        };
83        Ok(std::env::var("MACKEREL_PLUGIN_WORKDIR")
84            .map_or_else(
85                |_| std::env::temp_dir(),
86                |path| std::path::PathBuf::from(&path),
87            )
88            .join(name)
89            .to_str()
90            .ok_or("invalid plugin working directory")?
91            .to_owned())
92    }
93
94    #[doc(hidden)]
95    fn output_definitions(&self, out: &mut dyn std::io::Write) -> Result<(), String> {
96        writeln!(out, "# mackerel-agent-plugin").map_err(|e| format!("{}", e))?;
97        let prefix = self.metric_key_prefix();
98        let json = json!({
99            "graphs": self.graph_definition()
100                .iter()
101                .map(|graph|
102                    (
103                        if prefix.is_empty() {
104                            graph.name.clone()
105                        } else if graph.name.is_empty() {
106                            prefix.clone()
107                        } else {
108                            prefix.clone() + "." + graph.name.as_ref()
109                        },
110                        graph
111                    )
112                )
113                .collect::<HashMap<_, _>>(),
114        });
115        writeln!(out, "{}", json).map_err(|e| format!("{}", e))?;
116        Ok(())
117    }
118
119    fn run(&self) -> Result<(), String> {
120        let stdout = std::io::stdout();
121        let mut out = std::io::BufWriter::new(stdout.lock());
122        if std::env::var("MACKEREL_AGENT_PLUGIN_META").map_or(false, |value| !value.is_empty()) {
123            self.output_definitions(&mut out)
124        } else {
125            self.output_values(&mut out)
126        }
127    }
128}
129
130fn load_values(path: &str) -> Result<MetricValues, String> {
131    let file = std::fs::File::open(path).map_err(|e| format!("open {} failed: {}", path, e))?;
132    serde_json::de::from_reader(file).map_err(|e| format!("read {} failed: {}", path, e))
133}
134
135fn save_values(path: &str, metric_values: &MetricValues) -> Result<(), String> {
136    let bytes = serde_json::to_vec(metric_values).unwrap();
137    atomic_write(path, bytes.as_slice())
138}
139
140fn atomic_write(path: &str, bytes: &[u8]) -> Result<(), String> {
141    let tmp_path = &format!(
142        "{}.{}",
143        path,
144        std::time::SystemTime::now()
145            .duration_since(std::time::UNIX_EPOCH)
146            .map_err(|e| e.to_string())?
147            .as_secs_f64()
148    );
149    let mut file =
150        std::fs::File::create(tmp_path).map_err(|e| format!("open {} failed: {}", tmp_path, e))?;
151    file.write(bytes)
152        .map_err(|e| format!("write to {} failed: {}", tmp_path, e))?;
153    drop(file);
154    std::fs::rename(tmp_path, path).map_err(|e| {
155        let _ = std::fs::remove_file(tmp_path);
156        format!("rename {} to {} failed: {}", tmp_path, path, e)
157    })
158}
159
160fn format_values(
161    out: &mut dyn std::io::Write,
162    prefix: &str,
163    graph_name: &str,
164    metric: Metric,
165    metric_values: &MetricValues,
166    prev_metric_values: &MetricValues,
167) {
168    for (metric_name, value) in
169        collect_metric_values(graph_name, metric, metric_values, prev_metric_values)
170    {
171        if !value.is_nan() && value.is_finite() {
172            let name = if prefix.is_empty() {
173                metric_name
174            } else {
175                prefix.to_owned() + "." + metric_name.as_ref()
176            };
177            writeln!(out, "{}\t{}\t{}", name, value, metric_values.timestamp).unwrap();
178        }
179    }
180}
181
182#[auto_enum(Iterator)]
183fn collect_metric_values<'a>(
184    graph_name: &'a str,
185    metric: Metric,
186    metric_values: &'a MetricValues,
187    prev_metric_values: &'a MetricValues,
188) -> impl Iterator<Item = (String, f64)> + 'a {
189    let metric_name = if graph_name.is_empty() {
190        metric.name
191    } else {
192        graph_name.to_owned() + "." + &metric.name
193    };
194    let count = metric_name.chars().filter(|&c| c == '.').count();
195    if metric_name.contains('*') || metric_name.contains('#') {
196        metric_values
197            .values
198            .iter()
199            .filter(move |&(name, _)| {
200                name.chars().filter(|&c| c == '.').count() == count
201                    && metric_name.split('.').zip(name.split('.')).all(|(cs, ds)| {
202                        if cs == "*" || cs == "#" {
203                            !ds.is_empty()
204                                && ds.chars().all(
205                                    |c| matches!(c, 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_'),
206                                )
207                        } else {
208                            cs == ds
209                        }
210                    })
211            })
212            .filter_map(move |(metric_name, &value)| {
213                if metric.diff {
214                    prev_metric_values
215                        .values
216                        .get(metric_name)
217                        .and_then(|&prev_value| {
218                            calc_diff(
219                                value,
220                                metric_values.timestamp,
221                                prev_value,
222                                prev_metric_values.timestamp,
223                            )
224                        })
225                } else {
226                    Some(value)
227                }
228                .map(|value| (metric_name.clone(), value))
229            })
230    } else {
231        metric_values
232            .values
233            .get(&metric_name)
234            .and_then(|&value| {
235                if metric.diff {
236                    prev_metric_values
237                        .values
238                        .get(&metric_name)
239                        .and_then(|&prev_value| {
240                            calc_diff(
241                                value,
242                                metric_values.timestamp,
243                                prev_value,
244                                prev_metric_values.timestamp,
245                            )
246                        })
247                } else {
248                    Some(value)
249                }
250            })
251            .map(|value| (metric_name, value))
252            .into_iter()
253    }
254}
255
256#[inline]
257fn calc_diff(value: f64, timestamp: i64, prev_value: f64, prev_timestamp: i64) -> Option<f64> {
258    if prev_timestamp < timestamp - 600 || timestamp <= prev_timestamp || prev_value > value {
259        None
260    } else {
261        Some((value - prev_value) / ((timestamp - prev_timestamp) as f64 / 60.0))
262    }
263}
mackerel_plugin/plugin.rs

mackerel_plugin/
plugin.rs
