Skip to main content

tracing_config/tracing/
sifting_layer.rs

1use std::fmt::Debug;
2use std::collections::HashMap;
3use std::hash::Hash;
4use std::sync::RwLock;
5
6use serde_json::value::Value as JsonValue;
7
8use ::tracing as t;
9use ::tracing_subscriber as ts;
10
11use t::{Subscriber, Event as TEvent, Level as TLevel};
12use t::span::{Id as SpanId, Record as TSpanRecord, Attributes as TSpanAttributes};
13
14use ts::layer::{Context, Layer as TsLayer};
15use ts::registry::{SpanRef, LookupSpan};
16
17use super::SpanRecord;
18
19/// A Selector is a set of span metadata and span keys, it's purpose is to sift spans.
20#[derive(Debug, Clone, Default, Hash, PartialEq, Eq)]
21pub struct Selector {
22    level: bool,
23    name: bool,
24    target: bool,
25    module_path: bool,
26    file: bool,
27    line: bool,
28    fields: Vec<String>,
29}
30
31/// Given a [`Selector`][struct@Selector], `Values` contains the values of a certain span.
32#[derive(Debug, Clone, Hash, PartialEq, Eq)]
33pub struct Values {
34    pub level: Option<TLevel>,
35    pub name: Option<&'static str>,
36    pub target: Option<String>,
37    pub module_path: Option<String>,
38    pub file: Option<String>,
39    pub line: Option<u32>,
40    pub field_values: Vec<Option<String>>,
41}
42
43/// A [`tracing-subscriber`][mod@ts] [`Layer`][trait@TsLayer] that sifts trough spans and forwards events to dynamically created Layers.
44///
45/// Just like all the layers it's generic over `S` which is a `Subscriber`. It is also generic over `L` which is the dynamically created layer and `FnLB` (LB stand for Layer Builder) which is the closure responsible to create a new Layer at runtime.
46///
47/// The `FnLB` will receive a [`Selector`][struct@Selector] and a [`SelectorValues`][struct@Values] and is responsible to create a unique layer for the specific set of `SelectorValues`.
48/// `FnLB` will be called once for every distinct set of `SelectorValues`.
49pub struct Layer<S, L, FnLB>
50where
51    S: Subscriber,
52    S: for<'lookup> LookupSpan<'lookup>,
53    L: TsLayer<S> + Send + Sync + 'static,
54    FnLB: FnMut(&Selector, &Values) -> L,
55{
56    sift_selector: Selector,
57    layer_cache: RwLock<HashMap<Values, L>>,
58    layer_builder: RwLock<FnLB>,
59    _marker: std::marker::PhantomData<S>,
60}
61
62impl Selector {
63    pub fn new() -> Self {
64        Self {
65            level: false,
66            name: false,
67            target: false,
68            module_path: false,
69            file: false,
70            line: false,
71            fields: Vec::new(),
72        }
73    }
74
75    pub fn level(mut self) -> Self {
76        self.level = true;
77        self
78    }
79
80    pub fn name(mut self) -> Self {
81        self.name = true;
82        self
83    }
84
85    pub fn target(mut self) -> Self {
86        self.target = true;
87        self
88    }
89
90    pub fn module_path(mut self) -> Self {
91        self.module_path = true;
92        self
93    }
94
95    pub fn file(mut self) -> Self {
96        self.file = true;
97        self
98    }
99
100    pub fn line(mut self) -> Self {
101        self.line = true;
102        self
103    }
104
105    pub fn field(mut self, field: &str) -> Self {
106        self.fields.push(field.to_owned());
107        self
108    }
109
110    pub fn index_of(&self, key: &str) -> Option<usize> {
111        for (i, k) in self.fields.iter().enumerate() {
112            if key == k {
113                return Some(i);
114            }
115        }
116        None
117    }
118
119    /// Given an `input` string and [`SelectorValues`][struct@Values] will replace placeholders.
120    pub fn resolve_variable(&self, input: &str, values: &Values) -> String {
121        // static REGEX_STR: &str = r"\$\{sl\:([^}]+)\}";
122
123        use crate::interpolate::*;
124
125        resolve_infallible(input, |scheme, key| {
126            if scheme != "sl" {
127                return None;
128            }
129
130            let s = match key {
131                "meta:level" => values.level.map(|level| level.to_string()),
132                "meta:name" => values.name.map(|name| name.to_owned()),
133                "meta:target" => values.target.clone(),
134                "meta:module_path" => values.module_path.clone(),
135                "meta:file" => values.file.clone(),
136                "meta:line" => values.line.map(|line| format!("{line}")),
137                key => {
138                    if let Some(idx) = self.index_of(key) {
139                        if let Some(Some(val)) = values.field_values.get(idx) {
140                            Some(val.to_owned())
141                        } else {
142                            None
143                        }
144                    } else {
145                        None
146                    }
147                }
148            };
149
150            if let Some(s) = s {
151                Some(s)
152            } else {
153                Some("none".to_owned())
154            }
155        })
156    }
157}
158
159impl Values {
160    pub fn new<S>(sift_selector: &Selector, span: Option<SpanRef<'_, S>>) -> Self
161    where
162        S: Subscriber,
163        S: for<'lookup> LookupSpan<'lookup>,
164    {
165        let span = match span {
166            Some(span) => span,
167            None => {
168                return Self {
169                    level: None,
170                    name: None,
171                    target: None,
172                    module_path: None,
173                    file: None,
174                    line: None,
175                    field_values: vec![None; sift_selector.fields.len()],
176                };
177            }
178        };
179
180        let metadata = span.metadata();
181
182        let module_path = if sift_selector.module_path {
183            metadata
184                .module_path()
185                .map(|module_path| module_path.to_owned())
186        } else {
187            None
188        };
189
190        let file = if sift_selector.file {
191            metadata.file().map(|file| file.to_owned())
192        } else {
193            None
194        };
195
196        let level = if sift_selector.level {
197            Some(*metadata.level())
198        } else {
199            None
200        };
201        let name = if sift_selector.name {
202            Some(metadata.name())
203        } else {
204            None
205        };
206        let target = if sift_selector.target {
207            Some(metadata.target().to_owned())
208        } else {
209            None
210        };
211        let line = if sift_selector.line {
212            metadata.line()
213        } else {
214            None
215        };
216
217        let mut field_values: Vec<Option<String>> = vec![None; sift_selector.fields.len()];
218        get_field_values_recursive(&sift_selector.fields, &mut field_values, &span);
219
220        return Self {
221            level,
222            name,
223            target,
224            module_path,
225            file,
226            line,
227            field_values,
228        };
229
230        // private inner helper function
231        fn get_field_values_recursive<S>(
232            fields_keys: &Vec<String>,
233            fields_values: &mut Vec<Option<String>>,
234            span: &SpanRef<'_, S>,
235        ) where
236            S: Subscriber,
237            S: for<'lookup> LookupSpan<'lookup>,
238        {
239            let span_extensions = span.extensions();
240            if let Some(span_values) = span_extensions.get::<SpanRecord>() {
241                for (key_index, key) in fields_keys.iter().enumerate() {
242                    match fields_values.get(key_index) {
243                        // the None case never happens
244                        // while Some(Some(_)) means that a value already exists and there is no need to override it.
245                        None | Some(Some(_)) => continue,
246                        // space for the value exists but it's not yet set.
247                        Some(None) => {
248                            if let Some(json_value) = span_values.map.get(key as &str) {
249                                let value = json_value_to_string(json_value);
250                                fields_values.insert(key_index, Some(value));
251                            }
252                        }
253                    }
254                }
255                if let Some(span_parent) = span.parent() {
256                    get_field_values_recursive(fields_keys, fields_values, &span_parent);
257                }
258            }
259        }
260    }
261}
262
263impl<S, L, FnLB> Layer<S, L, FnLB>
264where
265    S: Subscriber,
266    S: for<'lookup> LookupSpan<'lookup>,
267    L: TsLayer<S> + Send + Sync + 'static,
268    FnLB: FnMut(&Selector, &Values) -> L,
269{
270    /// To create a `SiftingLayer` a selector must be provided as well as a closure that creates layers based on the selectors values.
271    pub fn new(sift_selector: Selector, layer_builder: FnLB) -> Self {
272        Self {
273            sift_selector,
274            layer_cache: RwLock::new(HashMap::new()),
275            layer_builder: RwLock::new(layer_builder),
276            _marker: std::marker::PhantomData {},
277        }
278    }
279
280    fn get_layer<F>(&self, ssv: Values, f: F)
281    where
282        F: FnOnce(&L),
283    {
284        {
285            let layer_cache = rw_lock_read(&self.layer_cache, "layer_cache");
286
287            if let Some(layer) = layer_cache.get(&ssv) {
288                f(layer);
289                return;
290            }
291        }
292
293        let mut layer_cache = rw_lock_write(&self.layer_cache, "layer_cache");
294
295        if let Some(layer) = layer_cache.get(&ssv) {
296            f(layer);
297            return;
298        }
299
300        let layer = self.build_layer(&ssv);
301
302        f(&layer);
303
304        layer_cache.insert(ssv, layer);
305    }
306
307    fn build_layer(&self, ssv: &Values) -> L {
308        let mut layer_builder = rw_lock_write(&self.layer_builder, "layer_builder");
309        layer_builder(&self.sift_selector, ssv)
310    }
311}
312
313impl<S, L, FnLB> TsLayer<S> for Layer<S, L, FnLB>
314where
315    S: Subscriber,
316    S: for<'lookup> LookupSpan<'lookup>,
317    L: TsLayer<S> + Send + Sync + 'static,
318    FnLB: FnMut(&Selector, &Values) -> L + 'static,
319{
320    fn on_new_span(&self, attrs: &TSpanAttributes<'_>, id: &SpanId, ctx: Context<'_, S>) {
321        let ssv = Values::new(&self.sift_selector, ctx.span(id));
322        self.get_layer(ssv, |layer| {
323            layer.on_new_span(attrs, id, ctx);
324        });
325    }
326
327    fn on_record(&self, id: &SpanId, values: &TSpanRecord<'_>, ctx: Context<'_, S>) {
328        let ssv = Values::new(&self.sift_selector, ctx.span(id));
329        self.get_layer(ssv, |layer| {
330            layer.on_record(id, values, ctx);
331        });
332    }
333
334    fn on_enter(&self, id: &SpanId, ctx: Context<'_, S>) {
335        let ssv = Values::new(&self.sift_selector, ctx.span(id));
336        self.get_layer(ssv, |layer| {
337            layer.on_enter(id, ctx);
338        });
339    }
340
341    fn on_exit(&self, id: &SpanId, ctx: Context<'_, S>) {
342        let ssv = Values::new(&self.sift_selector, ctx.span(id));
343        self.get_layer(ssv, |layer| {
344            layer.on_exit(id, ctx);
345        });
346    }
347
348    fn on_close(&self, id: SpanId, ctx: Context<'_, S>) {
349        let ssv = Values::new(&self.sift_selector, ctx.span(&id));
350        self.get_layer(ssv, |layer| {
351            layer.on_close(id, ctx);
352        });
353    }
354
355    fn on_event(&self, event: &TEvent<'_>, ctx: Context<'_, S>) {
356        let span = ctx.current_span();
357        let span = if let Some(id) = span.id() {
358            ctx.span(id)
359        } else {
360            None
361        };
362        let ssv = Values::new(&self.sift_selector, span);
363        self.get_layer(ssv, |layer| {
364            layer.on_event(event, ctx);
365        });
366    }
367}
368
369#[inline(always)]
370fn rw_lock_read<'a, T>(
371    rw_lock: &'a RwLock<T>,
372    lock_name: &str,
373) -> std::sync::RwLockReadGuard<'a, T> {
374    match rw_lock.read() {
375        Ok(v) => v,
376        Err(v_error) => {
377            print_error(lock_name, "read", "lock poisoned");
378            v_error.into_inner()
379        }
380    }
381}
382
383#[inline(always)]
384fn rw_lock_write<'a, T>(
385    rw_lock: &'a RwLock<T>,
386    lock_name: &str,
387) -> std::sync::RwLockWriteGuard<'a, T> {
388    match rw_lock.write() {
389        Ok(v) => v,
390        Err(v_error) => {
391            print_error(lock_name, "read", "lock poisoned");
392            v_error.into_inner()
393        }
394    }
395}
396
397#[inline(always)]
398fn print_error(name: &str, func: &str, msg: &str) {
399    let now = chrono::Local::now();
400    let error = format!("{now} -> ERROR : SiftingLayer -- {name}.{func}() {msg} !");
401    println!("{error}");
402    eprintln!("{error}");
403}
404
405#[inline(always)]
406fn json_value_to_string(value: &JsonValue) -> String {
407    match value {
408        JsonValue::String(s) => s.clone(),
409        _ => value.to_string(),
410    }
411}
412
413#[test]
414fn to_some_repl_test() {
415    let s = Selector::new().field("k").field("ve");
416
417    let v = Values {
418        level: Some(TLevel::INFO),
419        name: None,
420        target: None,
421        module_path: None,
422        file: None,
423        line: None,
424        field_values: {
425            let mut v = Vec::new();
426            v.push(Some("k_v".to_owned()));
427            v.push(None);
428            v
429        },
430    };
431
432    let k = "some good string `${sl:meta:level}` containing : `${sl:k}` ${sl:ve} and `${sl:k}` ${sl:k2} ${sl:k} then some".to_string();
433    let r = s.resolve_variable(&k, &v);
434
435    println!("v = {r}");
436}