msgpack_tracing/
tape.rs

1use chrono::{DateTime, Utc};
2use std::{
3    num::NonZeroU64,
4    ops::DerefMut,
5    sync::{Mutex, MutexGuard},
6};
7use tracing::{
8    Level, Subscriber,
9    field::{Field, Visit},
10    span,
11};
12use tracing_subscriber::{Layer, registry::LookupSpan};
13
14pub trait TapeMachine<I>: Send + 'static
15where
16    I: InstructionSetTrait,
17{
18    fn needs_restart(&mut self) -> bool;
19    fn handle(&mut self, instruction: I::Instruction<'_>);
20}
21
22pub trait InstructionSetTrait {
23    type Instruction<'a>: InstructionTrait;
24}
25pub struct InstructionSet;
26impl InstructionSetTrait for InstructionSet {
27    type Instruction<'a> = Instruction<'a>;
28}
29
30pub trait InstructionTrait: Copy {
31    fn id(self) -> InstructionId;
32}
33
34#[derive(Clone, Copy, Debug)]
35pub enum Instruction<'a> {
36    Restart,
37    NewSpan {
38        parent: Option<NonZeroU64>,
39        span: NonZeroU64,
40        name: &'a str,
41    },
42    FinishedSpan,
43    NewRecord(NonZeroU64),
44    FinishedRecord,
45    StartEvent {
46        time: DateTime<Utc>,
47        span: Option<NonZeroU64>,
48        target: &'a str,
49        priority: Level,
50    },
51    FinishedEvent,
52    AddValue(FieldValue<'a, &'a str>),
53    DeleteSpan(NonZeroU64),
54}
55impl InstructionTrait for Instruction<'_> {
56    fn id(self) -> InstructionId {
57        match self {
58            Instruction::Restart => InstructionId::Restart,
59            Instruction::NewSpan { .. } => InstructionId::NewSpan,
60            Instruction::FinishedSpan => InstructionId::FinishedSpan,
61            Instruction::NewRecord(..) => InstructionId::NewRecord,
62            Instruction::FinishedRecord => InstructionId::FinishedRecord,
63            Instruction::StartEvent { .. } => InstructionId::StartEvent,
64            Instruction::FinishedEvent => InstructionId::FinishedEvent,
65            Instruction::AddValue(..) => InstructionId::AddValue,
66            Instruction::DeleteSpan(..) => InstructionId::DeleteSpan,
67        }
68    }
69}
70
71#[derive(Clone, Copy, Debug)]
72pub enum InstructionId {
73    Restart,
74    NewString,
75    NewSpan,
76    FinishedSpan,
77    NewRecord,
78    FinishedRecord,
79    StartEvent,
80    FinishedEvent,
81    AddValue,
82    DeleteSpan,
83}
84impl From<InstructionId> for u8 {
85    fn from(val: InstructionId) -> Self {
86        match val {
87            InstructionId::Restart => 255,
88            InstructionId::NewString => 1,
89            InstructionId::NewSpan => 2,
90            InstructionId::FinishedSpan => 4,
91            InstructionId::NewRecord => 8,
92            InstructionId::FinishedRecord => 16,
93            InstructionId::StartEvent => 32,
94            InstructionId::FinishedEvent => 64,
95            InstructionId::AddValue => 128,
96            InstructionId::DeleteSpan => 0,
97        }
98    }
99}
100impl TryFrom<u8> for InstructionId {
101    type Error = u8;
102
103    fn try_from(value: u8) -> Result<Self, Self::Error> {
104        Ok(match value {
105            255 => InstructionId::Restart,
106            1 => InstructionId::NewString,
107            2 => InstructionId::NewSpan,
108            4 => InstructionId::FinishedSpan,
109            8 => InstructionId::NewRecord,
110            16 => InstructionId::FinishedRecord,
111            32 => InstructionId::StartEvent,
112            64 => InstructionId::FinishedEvent,
113            128 => InstructionId::AddValue,
114            0 => InstructionId::DeleteSpan,
115            e => return Err(e),
116        })
117    }
118}
119
120#[derive(Clone, Copy, Debug)]
121pub struct FieldValue<'a, S> {
122    pub name: S,
123    pub value: Value<'a, S>,
124}
125impl<'a> FieldValue<'a, &'a str> {
126    pub fn to_owned(self) -> FieldValueOwned {
127        FieldValueOwned {
128            name: self.name.to_owned(),
129            value: self.value.to_owned(),
130        }
131    }
132}
133
134#[derive(Clone)]
135pub struct FieldValueOwned {
136    pub name: String,
137    pub value: ValueOwned,
138}
139impl FieldValueOwned {
140    pub fn as_ref(&self) -> FieldValue<&str> {
141        FieldValue {
142            name: &self.name,
143            value: self.value.as_ref(),
144        }
145    }
146}
147
148#[derive(Clone, Copy, Debug)]
149pub enum Value<'a, S> {
150    Debug(S),
151    String(S),
152    Float(f64),
153    Integer(i64),
154    Unsigned(u64),
155    Bool(bool),
156    ByteArray(&'a [u8]),
157}
158impl<S> From<f64> for Value<'_, S> {
159    fn from(value: f64) -> Self {
160        Value::Float(value)
161    }
162}
163impl<S> From<i64> for Value<'_, S> {
164    fn from(value: i64) -> Self {
165        Self::Integer(value)
166    }
167}
168impl<S> From<u64> for Value<'_, S> {
169    fn from(value: u64) -> Self {
170        Self::Unsigned(value)
171    }
172}
173impl<S> From<bool> for Value<'_, S> {
174    fn from(value: bool) -> Self {
175        Value::Bool(value)
176    }
177}
178impl<'a, S> From<&'a [u8]> for Value<'a, S> {
179    fn from(value: &'a [u8]) -> Self {
180        Value::ByteArray(value)
181    }
182}
183impl<'a> Value<'a, &'a str> {
184    fn to_owned(self) -> ValueOwned {
185        match self {
186            Value::Debug(str) => ValueOwned::Debug(str.to_owned()),
187            Value::String(str) => ValueOwned::String(str.to_owned()),
188            Value::Float(value) => ValueOwned::Float(value),
189            Value::Integer(value) => ValueOwned::Integer(value),
190            Value::Unsigned(value) => ValueOwned::Unsigned(value),
191            Value::Bool(value) => ValueOwned::Bool(value),
192            Value::ByteArray(items) => ValueOwned::ByteArray(items.to_owned()),
193        }
194    }
195}
196
197#[derive(Clone, Debug)]
198pub enum ValueOwned {
199    Debug(String),
200    String(String),
201    Float(f64),
202    Integer(i64),
203    Unsigned(u64),
204    Bool(bool),
205    ByteArray(Vec<u8>),
206}
207impl ValueOwned {
208    pub fn as_ref(&self) -> Value<&str> {
209        match self {
210            ValueOwned::Debug(value) => Value::Debug(value),
211            ValueOwned::String(value) => Value::String(value),
212            ValueOwned::Float(value) => Value::Float(*value),
213            ValueOwned::Integer(value) => Value::Integer(*value),
214            ValueOwned::Unsigned(value) => Value::Unsigned(*value),
215            ValueOwned::Bool(value) => Value::Bool(*value),
216            ValueOwned::ByteArray(items) => Value::ByteArray(items),
217        }
218    }
219}
220
221pub struct TapeMachineLogger<T> {
222    inner: Mutex<TapeMachineLoggerInner<T>>,
223}
224impl<T> TapeMachineLogger<T>
225where
226    T: TapeMachine<InstructionSet>,
227{
228    pub fn new(mut machine: T) -> Self {
229        machine.handle(Instruction::Restart);
230        TapeMachineLogger {
231            inner: Mutex::new(TapeMachineLoggerInner { machine }),
232        }
233    }
234
235    fn machine(&self) -> MutexGuard<'_, TapeMachineLoggerInner<T>> {
236        let mut machine = self.inner.lock().unwrap();
237        if machine.machine.needs_restart() {
238            machine.handle(Instruction::Restart);
239        }
240        machine
241    }
242}
243impl<T, S> Layer<S> for TapeMachineLogger<T>
244where
245    T: TapeMachine<InstructionSet>,
246    S: Subscriber + for<'a> LookupSpan<'a>,
247{
248    fn on_new_span(
249        &self,
250        attrs: &span::Attributes<'_>,
251        id: &span::Id,
252        ctx: tracing_subscriber::layer::Context<'_, S>,
253    ) {
254        let mut machine = self.machine();
255        let name = attrs.metadata().name();
256        let span = ctx.span(id).unwrap();
257        machine.handle(Instruction::NewSpan {
258            parent: span.parent().map(|parent| parent.id().into_non_zero_u64()),
259            span: id.into_non_zero_u64(),
260            name,
261        });
262        attrs.record(&mut VisitMachine(machine.deref_mut()));
263        machine.handle(Instruction::FinishedSpan);
264    }
265
266    fn on_record(
267        &self,
268        id: &span::Id,
269        values: &span::Record<'_>,
270        _ctx: tracing_subscriber::layer::Context<'_, S>,
271    ) {
272        let mut machine = self.machine();
273        machine.handle(Instruction::NewRecord(id.into_non_zero_u64()));
274        values.record(&mut VisitMachine(machine.deref_mut()));
275        machine.handle(Instruction::FinishedRecord);
276    }
277
278    fn on_event(&self, event: &tracing::Event<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
279        let mut machine = self.machine();
280
281        let time = Utc::now();
282        let span = ctx
283            .event_span(event)
284            .map(|span| span.id().into_non_zero_u64());
285        let priority = *event.metadata().level();
286        let target = event.metadata().target();
287        machine.handle(Instruction::StartEvent {
288            time,
289            span,
290            target,
291            priority,
292        });
293        event.record(&mut VisitMachine(machine.deref_mut()));
294
295        machine.handle(Instruction::FinishedEvent);
296    }
297
298    fn on_close(&self, id: span::Id, _ctx: tracing_subscriber::layer::Context<'_, S>) {
299        let mut machine = self.machine();
300        machine.handle(Instruction::DeleteSpan(id.into_non_zero_u64()));
301    }
302}
303
304struct TapeMachineLoggerInner<T> {
305    machine: T,
306}
307impl<T> TapeMachineLoggerInner<T>
308where
309    T: TapeMachine<InstructionSet>,
310{
311    fn field_value<'a, V>(&mut self, field: &Field, value: V) -> FieldValue<'a, &'a str>
312    where
313        V: Into<Value<'a, &'a str>>,
314    {
315        let name = field.name();
316        let value = value.into();
317
318        FieldValue { name, value }
319    }
320
321    fn handle(&mut self, instruction: Instruction) {
322        self.machine.handle(instruction);
323    }
324}
325
326struct VisitMachine<'a, T>(&'a mut TapeMachineLoggerInner<T>);
327impl<T> Visit for VisitMachine<'_, T>
328where
329    T: TapeMachine<InstructionSet>,
330{
331    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
332        let value = format!("{value:?}");
333        let value = self.0.field_value(field, Value::Debug(value.as_str()));
334        self.0.handle(Instruction::AddValue(value));
335    }
336
337    fn record_f64(&mut self, field: &Field, value: f64) {
338        let value = self.0.field_value(field, value);
339        self.0.handle(Instruction::AddValue(value));
340    }
341
342    fn record_i64(&mut self, field: &Field, value: i64) {
343        let value = self.0.field_value(field, value);
344        self.0.handle(Instruction::AddValue(value));
345    }
346
347    fn record_u64(&mut self, field: &Field, value: u64) {
348        let value = self.0.field_value(field, value);
349        self.0.handle(Instruction::AddValue(value));
350    }
351
352    fn record_i128(&mut self, field: &Field, value: i128) {
353        self.record_bytes(field, &value.to_le_bytes());
354    }
355
356    fn record_u128(&mut self, field: &Field, value: u128) {
357        self.record_bytes(field, &value.to_le_bytes());
358    }
359
360    fn record_bool(&mut self, field: &Field, value: bool) {
361        let value = self.0.field_value(field, value);
362        self.0.handle(Instruction::AddValue(value));
363    }
364
365    fn record_str(&mut self, field: &Field, value: &str) {
366        let value = self.0.field_value(field, Value::String(value));
367        self.0.handle(Instruction::AddValue(value));
368    }
369
370    fn record_bytes(&mut self, field: &Field, value: &[u8]) {
371        let value = self.0.field_value(field, value);
372        self.0.handle(Instruction::AddValue(value));
373    }
374
375    fn record_error(&mut self, field: &Field, value: &(dyn std::error::Error + 'static)) {
376        self.record_debug(field, &value.to_string())
377    }
378}
379
380#[derive(Clone)]
381pub struct SpanRecords {
382    pub parent: Option<NonZeroU64>,
383    pub name: String,
384    pub records: Vec<FieldValueOwned>,
385}
386impl SpanRecords {
387    pub fn lost(span: NonZeroU64) -> Self {
388        Self {
389            parent: None,
390            name: format!("span-{span}"),
391            records: Default::default(),
392        }
393    }
394}