1use reifydb_core::profiler::ProfilerCategoryId;
5use serde::{Deserialize, Serialize};
6use tracing::{Level, level_filters::LevelFilter};
7
8#[repr(u8)]
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
10pub enum ProfilerCategory {
11 Query = 0,
12 Txn = 1,
13 Storage = 2,
14 Plan = 3,
15 Cdc = 4,
16 Flow = 5,
17}
18
19pub const ALL_CATEGORIES: [ProfilerCategory; 6] = [
20 ProfilerCategory::Query,
21 ProfilerCategory::Txn,
22 ProfilerCategory::Storage,
23 ProfilerCategory::Plan,
24 ProfilerCategory::Cdc,
25 ProfilerCategory::Flow,
26];
27
28impl ProfilerCategory {
29 pub const fn as_id(self) -> ProfilerCategoryId {
30 ProfilerCategoryId(self as u8)
31 }
32
33 pub const fn from_id(id: ProfilerCategoryId) -> Option<Self> {
34 match id.0 {
35 0 => Some(ProfilerCategory::Query),
36 1 => Some(ProfilerCategory::Txn),
37 2 => Some(ProfilerCategory::Storage),
38 3 => Some(ProfilerCategory::Plan),
39 4 => Some(ProfilerCategory::Cdc),
40 5 => Some(ProfilerCategory::Flow),
41 _ => None,
42 }
43 }
44
45 pub fn from_span_name(name: &str) -> Option<Self> {
46 if name.starts_with("flow::engine::") {
47 Some(ProfilerCategory::Flow)
48 } else if name.starts_with("transaction::") {
49 Some(ProfilerCategory::Txn)
50 } else if name.starts_with("store::single::")
51 || name.starts_with("store::multi::")
52 || name.starts_with("drop::")
53 {
54 Some(ProfilerCategory::Storage)
55 } else if name.starts_with("volcano::") || name.starts_with("vm::") {
56 Some(ProfilerCategory::Query)
57 } else if name.starts_with("rql::") || name.starts_with("catalog::") {
58 Some(ProfilerCategory::Plan)
59 } else if name.starts_with("cdc::") {
60 Some(ProfilerCategory::Cdc)
61 } else {
62 None
63 }
64 }
65
66 pub const fn name(self) -> &'static str {
67 match self {
68 ProfilerCategory::Query => "query",
69 ProfilerCategory::Txn => "txn",
70 ProfilerCategory::Storage => "storage",
71 ProfilerCategory::Plan => "plan",
72 ProfilerCategory::Cdc => "cdc",
73 ProfilerCategory::Flow => "flow",
74 }
75 }
76}
77
78#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
79pub enum ProfilerLevel {
80 Trace,
81 Debug,
82 Info,
83 Warn,
84 Error,
85}
86
87impl ProfilerLevel {
88 pub fn as_level_filter(self) -> LevelFilter {
89 match self {
90 Self::Trace => LevelFilter::TRACE,
91 Self::Debug => LevelFilter::DEBUG,
92 Self::Info => LevelFilter::INFO,
93 Self::Warn => LevelFilter::WARN,
94 Self::Error => LevelFilter::ERROR,
95 }
96 }
97
98 pub fn admits(self, level: &Level) -> bool {
99 *level <= self.as_level_filter()
100 }
101}
102
103#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
104pub struct CategorySet {
105 levels: [Option<ProfilerLevel>; 6],
106}
107
108impl CategorySet {
109 pub const fn empty() -> Self {
110 Self {
111 levels: [None; 6],
112 }
113 }
114
115 pub const fn all() -> Self {
116 Self::all_at(ProfilerLevel::Trace)
117 }
118
119 pub const fn all_at(level: ProfilerLevel) -> Self {
120 Self {
121 levels: [Some(level); 6],
122 }
123 }
124
125 pub fn contains(&self, c: ProfilerCategory) -> bool {
126 self.levels[c as usize].is_some()
127 }
128
129 pub fn level_for(&self, c: ProfilerCategory) -> Option<ProfilerLevel> {
130 self.levels[c as usize]
131 }
132
133 pub fn insert(&mut self, c: ProfilerCategory) {
134 self.insert_at(c, ProfilerLevel::Trace);
135 }
136
137 pub fn insert_at(&mut self, c: ProfilerCategory, level: ProfilerLevel) {
138 self.levels[c as usize] = Some(level);
139 }
140
141 pub fn remove(&mut self, c: ProfilerCategory) {
142 self.levels[c as usize] = None;
143 }
144
145 pub fn with(mut self, c: ProfilerCategory) -> Self {
146 self.insert(c);
147 self
148 }
149
150 pub fn with_level(mut self, c: ProfilerCategory, level: ProfilerLevel) -> Self {
151 self.insert_at(c, level);
152 self
153 }
154
155 pub fn without(mut self, c: ProfilerCategory) -> Self {
156 self.remove(c);
157 self
158 }
159
160 pub fn is_empty(&self) -> bool {
161 self.levels.iter().all(|l| l.is_none())
162 }
163}
164
165#[cfg(test)]
166mod tests {
167 use super::*;
168
169 #[test]
170 fn from_span_name_known_prefixes() {
171 assert_eq!(ProfilerCategory::from_span_name("flow::engine::apply"), Some(ProfilerCategory::Flow));
172 assert_eq!(
173 ProfilerCategory::from_span_name("flow::engine::process_batch"),
174 Some(ProfilerCategory::Flow)
175 );
176 assert_eq!(ProfilerCategory::from_span_name("transaction::commit"), Some(ProfilerCategory::Txn));
177 assert_eq!(ProfilerCategory::from_span_name("store::multi::write"), Some(ProfilerCategory::Storage));
178 assert_eq!(ProfilerCategory::from_span_name("store::single::scan"), Some(ProfilerCategory::Storage));
179 assert_eq!(ProfilerCategory::from_span_name("drop::range"), Some(ProfilerCategory::Storage));
180 assert_eq!(ProfilerCategory::from_span_name("volcano::project"), Some(ProfilerCategory::Query));
181 assert_eq!(ProfilerCategory::from_span_name("vm::executor"), Some(ProfilerCategory::Query));
182 assert_eq!(ProfilerCategory::from_span_name("rql::parse"), Some(ProfilerCategory::Plan));
183 assert_eq!(ProfilerCategory::from_span_name("catalog::lookup"), Some(ProfilerCategory::Plan));
184 assert_eq!(ProfilerCategory::from_span_name("cdc::write"), Some(ProfilerCategory::Cdc));
185 }
186
187 #[test]
188 fn from_span_name_unknown_returns_none() {
189 assert_eq!(ProfilerCategory::from_span_name("random::other::name"), None);
190 assert_eq!(ProfilerCategory::from_span_name(""), None);
191 assert_eq!(ProfilerCategory::from_span_name("tracing::subsystem::start"), None);
192 }
193
194 #[test]
195 fn id_roundtrip() {
196 for c in ALL_CATEGORIES {
197 assert_eq!(ProfilerCategory::from_id(c.as_id()), Some(c));
198 }
199 assert_eq!(ProfilerCategory::from_id(ProfilerCategoryId(6)), None);
200 assert_eq!(ProfilerCategory::from_id(ProfilerCategoryId(255)), None);
201 }
202
203 #[test]
204 fn category_set_membership() {
205 let mut s = CategorySet::empty();
206 assert!(s.is_empty());
207 s.insert(ProfilerCategory::Flow);
208 assert!(s.contains(ProfilerCategory::Flow));
209 assert!(!s.contains(ProfilerCategory::Query));
210 s.insert(ProfilerCategory::Query);
211 assert!(s.contains(ProfilerCategory::Query));
212 s.remove(ProfilerCategory::Flow);
213 assert!(!s.contains(ProfilerCategory::Flow));
214
215 let all = CategorySet::all();
216 for c in ALL_CATEGORIES {
217 assert!(all.contains(c));
218 }
219 }
220
221 #[test]
222 fn category_set_default_insert_is_trace() {
223 let mut s = CategorySet::empty();
224 s.insert(ProfilerCategory::Flow);
225 assert_eq!(s.level_for(ProfilerCategory::Flow), Some(ProfilerLevel::Trace));
226 }
227
228 #[test]
229 fn category_set_per_category_level() {
230 let mut s = CategorySet::empty();
231 s.insert_at(ProfilerCategory::Flow, ProfilerLevel::Trace);
232 s.insert_at(ProfilerCategory::Query, ProfilerLevel::Debug);
233 s.insert_at(ProfilerCategory::Storage, ProfilerLevel::Info);
234
235 assert_eq!(s.level_for(ProfilerCategory::Flow), Some(ProfilerLevel::Trace));
236 assert_eq!(s.level_for(ProfilerCategory::Query), Some(ProfilerLevel::Debug));
237 assert_eq!(s.level_for(ProfilerCategory::Storage), Some(ProfilerLevel::Info));
238 assert_eq!(s.level_for(ProfilerCategory::Plan), None);
239
240 assert!(s.contains(ProfilerCategory::Flow));
241 assert!(s.contains(ProfilerCategory::Query));
242 assert!(s.contains(ProfilerCategory::Storage));
243 assert!(!s.contains(ProfilerCategory::Plan));
244 }
245
246 #[test]
247 fn category_set_with_level_round_trips() {
248 let s = CategorySet::empty()
249 .with_level(ProfilerCategory::Flow, ProfilerLevel::Trace)
250 .with_level(ProfilerCategory::Plan, ProfilerLevel::Debug);
251 assert_eq!(s.level_for(ProfilerCategory::Flow), Some(ProfilerLevel::Trace));
252 assert_eq!(s.level_for(ProfilerCategory::Plan), Some(ProfilerLevel::Debug));
253 assert_eq!(s.level_for(ProfilerCategory::Cdc), None);
254 }
255
256 #[test]
257 fn category_set_all_at_sets_uniform_level() {
258 let s = CategorySet::all_at(ProfilerLevel::Debug);
259 for c in ALL_CATEGORIES {
260 assert_eq!(s.level_for(c), Some(ProfilerLevel::Debug));
261 }
262 }
263
264 #[test]
265 fn profile_level_admits_at_or_less_verbose() {
266 assert!(ProfilerLevel::Debug.admits(&Level::DEBUG));
267 assert!(ProfilerLevel::Debug.admits(&Level::INFO));
268 assert!(ProfilerLevel::Debug.admits(&Level::WARN));
269 assert!(ProfilerLevel::Debug.admits(&Level::ERROR));
270 assert!(!ProfilerLevel::Debug.admits(&Level::TRACE));
271
272 assert!(ProfilerLevel::Trace.admits(&Level::TRACE));
273 assert!(!ProfilerLevel::Error.admits(&Level::WARN));
274 }
275}