msgpack_tracing/
restart.rs

1use crate::tape::{Instruction, InstructionSet, SpanRecords, TapeMachine};
2use std::{collections::HashMap, num::NonZeroU64};
3
4pub struct RestartableMachine<T> {
5    forward: T,
6    span: HashMap<NonZeroU64, SpanRecords>,
7    current_span: Option<(NonZeroU64, SpanRecords)>,
8}
9impl<T> RestartableMachine<T>
10where
11    T: TapeMachine<InstructionSet>,
12{
13    pub fn new(forward: T) -> Self {
14        Self {
15            forward,
16            span: Default::default(),
17            current_span: None,
18        }
19    }
20}
21impl<T> TapeMachine<InstructionSet> for RestartableMachine<T>
22where
23    T: TapeMachine<InstructionSet>,
24{
25    fn needs_restart(&mut self) -> bool {
26        self.forward.needs_restart()
27    }
28
29    fn handle(&mut self, instruction: Instruction) {
30        match instruction {
31            Instruction::Restart => {
32                self.forward.handle(Instruction::Restart);
33
34                for (span, records) in self.span.iter() {
35                    self.forward.handle(Instruction::NewSpan {
36                        parent: records.parent,
37                        span: *span,
38                        name: records.name.as_ref(),
39                    });
40
41                    for record in records.records.iter() {
42                        self.forward.handle(Instruction::AddValue(record.as_ref()));
43                    }
44
45                    self.forward.handle(Instruction::FinishedSpan);
46                }
47            }
48            Instruction::NewSpan { parent, span, name } => {
49                assert!(self.current_span.is_none());
50                self.current_span = Some((
51                    span,
52                    SpanRecords {
53                        parent,
54                        name: name.to_owned(),
55                        records: Default::default(),
56                    },
57                ));
58
59                self.forward
60                    .handle(Instruction::NewSpan { parent, span, name });
61            }
62            Instruction::FinishedSpan => {
63                let (k, v) = self.current_span.take().unwrap();
64                self.span.insert(k, v);
65                self.forward.handle(Instruction::FinishedSpan)
66            }
67            Instruction::NewRecord(span) => {
68                assert!(self.current_span.is_none());
69                self.current_span = Some(self.span.remove_entry(&span).unwrap());
70                self.forward.handle(Instruction::NewRecord(span));
71            }
72            Instruction::FinishedRecord => {
73                let (k, v) = self.current_span.take().unwrap();
74                self.span.insert(k, v);
75                self.forward.handle(Instruction::FinishedRecord)
76            }
77            Instruction::StartEvent {
78                time,
79                span,
80                target,
81                priority,
82            } => {
83                self.forward.handle(Instruction::StartEvent {
84                    time,
85                    span,
86                    target,
87                    priority,
88                });
89            }
90            Instruction::FinishedEvent => self.forward.handle(Instruction::FinishedEvent),
91            Instruction::AddValue(field_value) => {
92                if let Some((_, current_span)) = self.current_span.as_mut() {
93                    current_span.records.push(field_value.to_owned());
94                }
95                self.forward.handle(Instruction::AddValue(field_value));
96            }
97            Instruction::DeleteSpan(span) => {
98                self.span.remove(&span);
99                self.forward.handle(Instruction::DeleteSpan(span));
100            }
101        }
102    }
103}
msgpack_tracing/restart.rs

msgpack_tracing/
restart.rs
