conjure_oxide/
rule_trace_aggregates.rs

1use std::collections::BTreeMap;
2use std::fs::{self, File};
3use std::io::Write as _;
4use std::path::{Path, PathBuf};
5use std::sync::{Arc, Mutex};
6
7use anyhow::Context as _;
8use tracing::field::{Field, Visit};
9use tracing::{Event, Subscriber};
10use tracing_subscriber::Layer;
11use tracing_subscriber::layer::Context;
12
13#[derive(Clone)]
14pub struct RuleTraceAggregatesHandle {
15    state: Arc<Mutex<RuleTraceAggregatesState>>,
16}
17
18pub struct RuleTraceAggregatesLayer {
19    state: Arc<Mutex<RuleTraceAggregatesState>>,
20}
21
22struct RuleTraceAggregatesState {
23    path: PathBuf,
24    tmp_path: PathBuf,
25    total_rule_applications: usize,
26    counts: BTreeMap<String, usize>,
27}
28
29#[derive(Default)]
30struct RuleNameVisitor {
31    rule_name: Option<String>,
32}
33
34impl RuleTraceAggregatesHandle {
35    pub fn new(path: PathBuf) -> anyhow::Result<Self> {
36        let state = RuleTraceAggregatesState::new(path);
37        state.write_snapshot()?;
38
39        Ok(Self {
40            state: Arc::new(Mutex::new(state)),
41        })
42    }
43
44    pub fn layer(&self) -> RuleTraceAggregatesLayer {
45        RuleTraceAggregatesLayer {
46            state: Arc::clone(&self.state),
47        }
48    }
49
50    pub fn flush(&self) {
51        self.state
52            .lock()
53            .expect("rule trace aggregate state lock poisoned")
54            .write_snapshot()
55            .expect("failed to write rule trace aggregates")
56    }
57}
58
59impl<S> Layer<S> for RuleTraceAggregatesLayer
60where
61    S: Subscriber,
62{
63    fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
64        let mut visitor = RuleNameVisitor::default();
65        event.record(&mut visitor);
66
67        let Some(rule_name) = visitor.rule_name else {
68            return;
69        };
70
71        self.state
72            .lock()
73            .expect("rule trace aggregate state lock poisoned")
74            .record_rule(rule_name)
75            .expect("failed to write rule trace aggregates");
76    }
77}
78
79impl RuleTraceAggregatesState {
80    fn new(path: PathBuf) -> Self {
81        Self {
82            tmp_path: temporary_output_path(&path),
83            path,
84            total_rule_applications: 0,
85            counts: BTreeMap::new(),
86        }
87    }
88
89    fn record_rule(&mut self, rule_name: String) -> anyhow::Result<()> {
90        self.total_rule_applications += 1;
91        *self.counts.entry(rule_name).or_insert(0) += 1;
92        self.write_snapshot()
93    }
94
95    fn write_snapshot(&self) -> anyhow::Result<()> {
96        let mut rows: Vec<_> = self.counts.iter().collect();
97        rows.sort_by(|(rule_name_a, count_a), (rule_name_b, count_b)| {
98            count_b
99                .cmp(count_a)
100                .then_with(|| rule_name_a.cmp(rule_name_b))
101        });
102
103        let mut file = File::create(&self.tmp_path).with_context(|| {
104            format!(
105                "Unable to create temporary aggregate trace file {}",
106                self.tmp_path.display()
107            )
108        })?;
109
110        writeln!(
111            file,
112            "total_rule_applications: {}",
113            self.total_rule_applications
114        )?;
115        for (rule_name, count) in rows {
116            writeln!(file, "{count:6} {rule_name}")?;
117        }
118        file.flush()
119            .expect("failed to flush temporary aggregate trace file");
120
121        fs::rename(&self.tmp_path, &self.path).with_context(|| {
122            format!(
123                "Unable to move aggregate trace file into place at {}",
124                self.path.display()
125            )
126        })?;
127
128        Ok(())
129    }
130}
131
132impl Visit for RuleNameVisitor {
133    fn record_str(&mut self, field: &Field, value: &str) {
134        if field.name() == "rule_name" {
135            self.rule_name = Some(value.to_owned());
136        }
137    }
138
139    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
140        if field.name() == "rule_name" && self.rule_name.is_none() {
141            self.rule_name = Some(format!("{value:?}").trim_matches('"').to_owned());
142        }
143    }
144}
145
146fn temporary_output_path(path: &Path) -> PathBuf {
147    let filename = path
148        .file_name()
149        .map(|name| name.to_string_lossy().into_owned())
150        .unwrap_or_else(|| "rule-trace-aggregates".to_owned());
151
152    path.with_file_name(format!(".{filename}.tmp"))
153}
conjure_oxide/rule_trace_aggregates.rs

conjure_oxide/
rule_trace_aggregates.rs
