1use std::{
2 cmp::{max, min},
3 env,
4 num::NonZeroUsize,
5 sync::{OnceLock, atomic::AtomicU64},
6};
7
8use rustc_hash::FxHashSet;
9
10use crate::{
11 self_time_tree::SelfTimeTree,
12 span::{Span, SpanEvent, SpanIndex},
13 span_ref::SpanRef,
14 timestamp::Timestamp,
15};
16
17pub type SpanId = NonZeroUsize;
18
19const CUT_OFF_DEPTH: u32 = 150;
20
21pub struct Store {
22 pub(crate) spans: Vec<Span>,
23 pub(crate) self_time_tree: Option<SelfTimeTree<SpanIndex>>,
24 max_self_time_lookup_time: AtomicU64,
25}
26
27fn new_root_span() -> Span {
28 Span {
29 parent: None,
30 depth: 0,
31 start: Timestamp::MAX,
32 category: "".into(),
33 name: "(root)".into(),
34 args: vec![],
35 events: vec![],
36 is_complete: true,
37 max_depth: OnceLock::new(),
38 self_allocations: 0,
39 self_allocation_count: 0,
40 self_deallocations: 0,
41 self_deallocation_count: 0,
42 total_allocations: OnceLock::new(),
43 total_deallocations: OnceLock::new(),
44 total_persistent_allocations: OnceLock::new(),
45 total_allocation_count: OnceLock::new(),
46 total_span_count: OnceLock::new(),
47 time_data: OnceLock::new(),
48 extra: OnceLock::new(),
49 names: OnceLock::new(),
50 }
51}
52
53impl Store {
54 pub fn new() -> Self {
55 Self {
56 spans: vec![new_root_span()],
57 self_time_tree: env::var("NO_CORRECTED_TIME")
58 .ok()
59 .is_none()
60 .then(SelfTimeTree::new),
61 max_self_time_lookup_time: AtomicU64::new(0),
62 }
63 }
64
65 pub fn reset(&mut self) {
66 self.spans.truncate(1);
67 self.spans[0] = new_root_span();
68 if let Some(tree) = self.self_time_tree.as_mut() {
69 *tree = SelfTimeTree::new();
70 }
71 *self.max_self_time_lookup_time.get_mut() = 0;
72 }
73
74 pub fn has_time_info(&self) -> bool {
75 self.self_time_tree
76 .as_ref()
77 .is_none_or(|tree| tree.len() > 0)
78 }
79
80 pub fn add_span(
81 &mut self,
82 parent: Option<SpanIndex>,
83 start: Timestamp,
84 category: String,
85 name: String,
86 args: Vec<(String, String)>,
87 outdated_spans: &mut FxHashSet<SpanIndex>,
88 ) -> SpanIndex {
89 let id = SpanIndex::new(self.spans.len()).unwrap();
90 self.spans.push(Span {
91 parent,
92 depth: 0,
93 start,
94 category,
95 name,
96 args,
97 events: vec![],
98 is_complete: false,
99 max_depth: OnceLock::new(),
100 self_allocations: 0,
101 self_allocation_count: 0,
102 self_deallocations: 0,
103 self_deallocation_count: 0,
104 total_allocations: OnceLock::new(),
105 total_deallocations: OnceLock::new(),
106 total_persistent_allocations: OnceLock::new(),
107 total_allocation_count: OnceLock::new(),
108 total_span_count: OnceLock::new(),
109 time_data: OnceLock::new(),
110 extra: OnceLock::new(),
111 names: OnceLock::new(),
112 });
113 let parent = if let Some(parent) = parent {
114 outdated_spans.insert(parent);
115 &mut self.spans[parent.get()]
116 } else {
117 &mut self.spans[0]
118 };
119 parent.start = min(parent.start, start);
120 let depth = parent.depth + 1;
121 if depth < CUT_OFF_DEPTH {
122 parent.events.push(SpanEvent::Child { index: id });
123 }
124 let span = &mut self.spans[id.get()];
125 span.depth = depth;
126 id
127 }
128
129 pub fn add_args(
130 &mut self,
131 span_index: SpanIndex,
132 args: Vec<(String, String)>,
133 outdated_spans: &mut FxHashSet<SpanIndex>,
134 ) {
135 let span = &mut self.spans[span_index.get()];
136 span.args.extend(args);
137 outdated_spans.insert(span_index);
138 }
139
140 pub fn set_max_self_time_lookup(&self, time: Timestamp) {
141 let time = *time;
142 let mut old = self
143 .max_self_time_lookup_time
144 .load(std::sync::atomic::Ordering::Relaxed);
145 while old < time {
146 match self.max_self_time_lookup_time.compare_exchange(
147 old,
148 time,
149 std::sync::atomic::Ordering::Relaxed,
150 std::sync::atomic::Ordering::Relaxed,
151 ) {
152 Ok(_) => break,
153 Err(real_old) => old = real_old,
154 }
155 }
156 }
157
158 fn insert_self_time(
159 &mut self,
160 start: Timestamp,
161 end: Timestamp,
162 span_index: SpanIndex,
163 outdated_spans: &mut FxHashSet<SpanIndex>,
164 ) {
165 if let Some(tree) = self.self_time_tree.as_mut() {
166 if Timestamp::from_value(*self.max_self_time_lookup_time.get_mut()) >= start {
167 tree.for_each_in_range(start, end, |_, _, span| {
168 outdated_spans.insert(*span);
169 });
170 }
171 tree.insert(start, end, span_index);
172 }
173 }
174
175 pub fn add_self_time(
176 &mut self,
177 span_index: SpanIndex,
178 start: Timestamp,
179 end: Timestamp,
180 outdated_spans: &mut FxHashSet<SpanIndex>,
181 ) {
182 let span = &mut self.spans[span_index.get()];
183 let time_data = span.time_data_mut();
184 if time_data.ignore_self_time {
185 return;
186 }
187 outdated_spans.insert(span_index);
188 time_data.self_time += end - start;
189 time_data.self_end = max(time_data.self_end, end);
190 span.events.push(SpanEvent::SelfTime { start, end });
191 self.insert_self_time(start, end, span_index, outdated_spans);
192 }
193
194 pub fn set_total_time(
195 &mut self,
196 span_index: SpanIndex,
197 start_time: Timestamp,
198 total_time: Timestamp,
199 outdated_spans: &mut FxHashSet<SpanIndex>,
200 ) {
201 let span = SpanRef {
202 span: &self.spans[span_index.get()],
203 store: self,
204 index: span_index.get(),
205 };
206 let mut children = span
207 .children()
208 .map(|c| (c.span.start, c.span.time_data().self_end, c.index()))
209 .collect::<Vec<_>>();
210 children.sort();
211 let self_end = start_time + total_time;
212 let mut self_time = Timestamp::ZERO;
213 let mut current = start_time;
214 let mut events = Vec::new();
215 for (start, end, index) in children {
216 if start > current {
217 if start > self_end {
218 events.push(SpanEvent::SelfTime {
219 start: current,
220 end: self_end,
221 });
222 self.insert_self_time(current, self_end, span_index, outdated_spans);
223 self_time += self_end - current;
224 break;
225 }
226 events.push(SpanEvent::SelfTime {
227 start: current,
228 end: start,
229 });
230 self.insert_self_time(current, start, span_index, outdated_spans);
231 self_time += start - current;
232 }
233 events.push(SpanEvent::Child { index });
234 current = max(current, end);
235 }
236 current -= start_time;
237 if current < total_time {
238 self_time += total_time - current;
239 events.push(SpanEvent::SelfTime {
240 start: current + start_time,
241 end: start_time + total_time,
242 });
243 self.insert_self_time(
244 current + start_time,
245 start_time + total_time,
246 span_index,
247 outdated_spans,
248 );
249 }
250 let span = &mut self.spans[span_index.get()];
251 outdated_spans.insert(span_index);
252 let time_data = span.time_data_mut();
253 time_data.self_time = self_time;
254 time_data.self_end = self_end;
255 span.events = events;
256 span.start = start_time;
257 }
258
259 pub fn set_parent(
260 &mut self,
261 span_index: SpanIndex,
262 parent: SpanIndex,
263 outdated_spans: &mut FxHashSet<SpanIndex>,
264 ) {
265 outdated_spans.insert(span_index);
266 let span = &mut self.spans[span_index.get()];
267
268 let old_parent = span.parent.replace(parent);
269 let old_parent = if let Some(parent) = old_parent {
270 outdated_spans.insert(parent);
271 &mut self.spans[parent.get()]
272 } else {
273 &mut self.spans[0]
274 };
275 if let Some(index) = old_parent
276 .events
277 .iter()
278 .position(|event| *event == SpanEvent::Child { index: span_index })
279 {
280 old_parent.events.remove(index);
281 }
282
283 outdated_spans.insert(parent);
284 let parent = &mut self.spans[parent.get()];
285 parent.events.push(SpanEvent::Child { index: span_index });
286 }
287
288 pub fn add_allocation(
289 &mut self,
290 span_index: SpanIndex,
291 allocation: u64,
292 count: u64,
293 outdated_spans: &mut FxHashSet<SpanIndex>,
294 ) {
295 let span = &mut self.spans[span_index.get()];
296 outdated_spans.insert(span_index);
297 span.self_allocations += allocation;
298 span.self_allocation_count += count;
299 }
300
301 pub fn add_deallocation(
302 &mut self,
303 span_index: SpanIndex,
304 deallocation: u64,
305 count: u64,
306 outdated_spans: &mut FxHashSet<SpanIndex>,
307 ) {
308 let span = &mut self.spans[span_index.get()];
309 outdated_spans.insert(span_index);
310 span.self_deallocations += deallocation;
311 span.self_deallocation_count += count;
312 }
313
314 pub fn complete_span(&mut self, span_index: SpanIndex) {
315 let span = &mut self.spans[span_index.get()];
316 span.is_complete = true;
317 }
318
319 pub fn invalidate_outdated_spans(&mut self, outdated_spans: &FxHashSet<SpanId>) {
320 fn invalidate_span(span: &mut Span) {
321 if let Some(time_data) = span.time_data.get_mut() {
322 time_data.end.take();
323 time_data.total_time.take();
324 time_data.corrected_self_time.take();
325 time_data.corrected_total_time.take();
326 }
327 span.total_allocations.take();
328 span.total_deallocations.take();
329 span.total_persistent_allocations.take();
330 span.total_allocation_count.take();
331 span.total_span_count.take();
332 span.extra.take();
333 }
334
335 for id in outdated_spans.iter() {
336 let mut span = &mut self.spans[id.get()];
337 loop {
338 invalidate_span(span);
339 let Some(parent) = span.parent else {
340 break;
341 };
342 if outdated_spans.contains(&parent) {
343 break;
344 }
345 span = &mut self.spans[parent.get()];
346 }
347 }
348
349 invalidate_span(&mut self.spans[0]);
350 }
351
352 pub fn root_spans(&self) -> impl Iterator<Item = SpanRef<'_>> {
353 self.spans[0].events.iter().filter_map(|event| match event {
354 &SpanEvent::Child { index: id } => Some(SpanRef {
355 span: &self.spans[id.get()],
356 store: self,
357 index: id.get(),
358 }),
359 _ => None,
360 })
361 }
362
363 pub fn root_span(&self) -> SpanRef<'_> {
364 SpanRef {
365 span: &self.spans[0],
366 store: self,
367 index: 0,
368 }
369 }
370
371 pub fn span(&self, id: SpanId) -> Option<(SpanRef<'_>, bool)> {
372 let id = id.get();
373 let is_graph = id & 1 == 1;
374 let index = id >> 1;
375 self.spans.get(index).map(|span| {
376 (
377 SpanRef {
378 span,
379 store: self,
380 index,
381 },
382 is_graph,
383 )
384 })
385 }
386}