pyroscope/backend/
ruleset.rs

1use super::{StackTrace, Tag};
2use crate::error::Result;
3use std::collections::hash_map::DefaultHasher;
4use std::collections::HashSet;
5use std::hash::Hasher;
6use std::sync::{Arc, Mutex};
7
8/// Profiling Rule
9#[derive(Debug, Eq, PartialEq, Hash, Clone)]
10pub enum Rule {
11    /// Global Tag
12    GlobalTag(Tag),
13    /// Thread Tag
14    ThreadTag(u64, Tag),
15}
16
17/// Ruleset is a set of rules that can be applied to a stacktrace. The rules
18/// are held in a Vector behind an Arc, so that they can be shared between
19/// threads.
20#[derive(Debug, Default, Clone)]
21pub struct Ruleset {
22    /// Rules vector
23    pub rules: Arc<Mutex<HashSet<Rule>>>,
24}
25
26impl Ruleset {
27    /// Create a new empty ruleset
28    pub fn new() -> Self {
29        Self {
30            rules: Arc::new(Mutex::new(HashSet::new())),
31        }
32    }
33
34    /// Add a rule to the ruleset
35    pub fn add_rule(&self, rule: Rule) -> Result<bool> {
36        let rules = self.rules.clone();
37
38        // Add the rule to the Ruleset
39        let insert = rules.lock()?.insert(rule);
40
41        Ok(insert)
42    }
43
44    /// Remove a rule from the ruleset
45    pub fn remove_rule(&self, rule: Rule) -> Result<bool> {
46        let rules = self.rules.clone();
47
48        // Remove the rule from the Ruleset
49        let remove = rules.lock()?.remove(&rule);
50
51        Ok(remove)
52    }
53
54    /// Return a list of all global tags
55    pub fn get_global_tags(&self) -> Result<Vec<Tag>> {
56        let rules = self.rules.clone();
57
58        let tags = rules
59            .lock()?
60            .iter()
61            .filter_map(|rule| match rule {
62                Rule::GlobalTag(tag) => Some(tag.to_owned()),
63                _ => None,
64            })
65            .collect();
66
67        Ok(tags)
68    }
69}
70
71impl std::ops::Add<&Ruleset> for StackTrace {
72    type Output = Self;
73    fn add(self, other: &Ruleset) -> Self {
74        // Get global Tags
75        let global_tags: Vec<Tag> = other.get_global_tags().unwrap_or_default();
76
77        // Filter Thread Tags
78        let stack_tags: Vec<Tag> = other
79            .rules
80            .lock()
81            .unwrap()
82            .iter()
83            .filter_map(|rule| {
84                if let Rule::ThreadTag(thread_id, tag) = rule {
85                    if let Some(stack_thread_id) = self.thread_id {
86                        // No PID, only thread ID to match
87                        if thread_id == &stack_thread_id {
88                            return Some(tag.clone());
89                        }
90                        if let (Some(stack_thread_id), Some(stack_pid)) = (self.thread_id, self.pid)
91                        {
92                            let mut hasher = DefaultHasher::new();
93                            hasher.write_u64(stack_thread_id % stack_pid as u64);
94                            let id = hasher.finish();
95                            if &id == thread_id {
96                                return Some(tag.clone());
97                            }
98                        }
99                    }
100                }
101                None
102            })
103            .collect();
104
105        // Add tags to metadata
106        let mut metadata = self.metadata.clone();
107        for tag in global_tags.iter().chain(stack_tags.iter()) {
108            metadata.add_tag(tag.clone());
109        }
110
111        Self {
112            pid: self.pid,
113            thread_id: self.thread_id,
114            thread_name: self.thread_name,
115            frames: self.frames,
116            metadata,
117        }
118    }
119}
pyroscope/backend/ruleset.rs

pyroscope/backend/
ruleset.rs
