1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

use crate::fields::{FieldConfig, FieldRecorder, FieldVisitor, RecordedValue, TryForEachField};
use std::sync::Arc;
use tracing_core::field::Field;
use tracing_core::span::{Attributes, Record};

pub trait SpanRecorder {
    fn record_span(&mut self, attrs: &Attributes<'_>);
    fn merge(&mut self, record: &Record<'_>);
}

pub struct DefaultSpanRecorder {
    config: Arc<FieldConfig>,
    fields: Vec<RecordedValue>,
}

impl SpanRecorder for DefaultSpanRecorder {
    fn record_span(&mut self, attrs: &Attributes<'_>) {
        attrs.record(&mut FieldVisitor::new(self))
    }

    fn merge(&mut self, record: &Record<'_>) {
        record.record(&mut FieldVisitor::new(self))
    }
}

impl TryForEachField for DefaultSpanRecorder {
    fn try_for_each<E, F: FnMut(&'static str, &RecordedValue) -> Result<(), E>>(
        &self,
        mut f: F,
    ) -> Result<(), E> {
        for (name, value) in self.config.span_field_names.iter().zip(self.fields.iter()) {
            f(name, value)?;
        }
        Ok(())
    }
}

impl FieldRecorder for DefaultSpanRecorder {
    fn record_field(&mut self, field: &Field, value: impl Into<RecordedValue>) {
        if let Some(i) = self.config.field_index(field) {
            self.fields[i] = value.into();
        }
    }
}

impl DefaultSpanRecorder {
    pub fn from_config(config: Arc<FieldConfig>) -> Self {
        let n = config.span_field_index.len();
        Self {
            config,
            fields: vec![RecordedValue::Unset; n],
        }
    }
}