Skip to main content

fibre_tracing/
lib.rs

1//! Span collection for the `tracing` ecosystem.
2//!
3//! `tracing` is fire and forget: a span goes to the subscriber and nothing
4//! comes back. That is right for logs and wrong for the three things that want
5//! to read a unit of work whole: an error page showing what this request did
6//! before it failed, a test asserting what a body actually called, and a
7//! header reporting per-step cost with no collector running.
8//!
9//! This crate is a layer that keeps them. A span marked as a **root** opens a
10//! trace, every span under it joins, and closing the root finishes the trace,
11//! offers it to whoever is listening and pushes it into a ring of recent ones.
12//!
13//! It knows nothing about requests. A web server calls its root a request, a
14//! queue worker calls it a job, a build calls it a compilation.
15//!
16//! ```ignore
17//! let (layer, traces) = fibre_tracing::layer();
18//! tracing_subscriber::registry().with(fibre_logging_layer).with(layer).init();
19//!
20//! #[tracing::instrument(fields(fibre.root = true))]
21//! fn handle() { ... }
22//!
23//! traces.current();     // inside the root, what has happened so far
24//! traces.recent(50);    // the last finished ones
25//! ```
26
27use std::collections::HashMap;
28use std::fmt;
29use std::sync::Arc;
30use std::sync::atomic::{AtomicU64, Ordering};
31use std::time::{Duration, Instant};
32
33use parking_lot::Mutex;
34use tracing::field::{Field, Visit};
35use tracing::span::{Attributes, Id, Record};
36use tracing::{Level, Subscriber};
37use tracing_subscriber::layer::Context;
38use tracing_subscriber::registry::LookupSpan;
39
40/// The field that marks a span as the top of a trace. Anything under it joins.
41pub const ROOT_FIELD: &str = "fibre.root";
42
43/// The field naming the outcome a span finished with, when it says.
44pub const OUTCOME_FIELD: &str = "fibre.outcome";
45
46/// How many finished traces the ring keeps by default.
47pub const DEFAULT_RING: usize = 256;
48
49/// A field value, in the shapes `tracing` records.
50#[derive(Debug, Clone, PartialEq)]
51pub enum Value {
52  Str(String),
53  Int(i64),
54  UInt(u64),
55  Float(f64),
56  Bool(bool),
57}
58
59impl fmt::Display for Value {
60  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61    match self {
62      Value::Str(v) => f.write_str(v),
63      Value::Int(v) => write!(f, "{v}"),
64      Value::UInt(v) => write!(f, "{v}"),
65      Value::Float(v) => write!(f, "{v}"),
66      Value::Bool(v) => write!(f, "{v}"),
67    }
68  }
69}
70
71/// One span of a trace, as it finished.
72#[derive(Debug, Clone)]
73pub struct Span {
74  pub name: &'static str,
75  pub target: String,
76  pub level: Level,
77  /// Its parent's index in the trace's own list; `None` at the root.
78  pub parent: Option<usize>,
79  /// How deep it sits, so a caller can print a waterfall without walking.
80  pub depth: usize,
81  /// When it opened, from the trace's own start.
82  pub at: Duration,
83  /// How long it was open. Zero while it still is.
84  pub duration: Duration,
85  pub fields: HashMap<String, Value>,
86}
87
88impl Span {
89  /// What `fibre.outcome` said, when it said anything.
90  pub fn outcome(&self) -> Option<&str> {
91    match self.fields.get(OUTCOME_FIELD) {
92      Some(Value::Str(s)) => Some(s.as_str()),
93      _ => None,
94    }
95  }
96}
97
98/// One unit of work: the root and everything that happened inside it, in the
99/// order the spans opened.
100#[derive(Debug, Clone)]
101pub struct Trace {
102  pub id: u64,
103  pub spans: Vec<Span>,
104  /// How long the root was open. Zero while the trace is still running.
105  pub duration: Duration,
106}
107
108impl Trace {
109  pub fn root(&self) -> Option<&Span> {
110    self.spans.first()
111  }
112
113  /// Every span of a name, which is how a caller asks "what calls did this
114  /// make" without knowing the shape of the tree.
115  pub fn named<'a>(&'a self, name: &'a str) -> impl Iterator<Item = &'a Span> + 'a {
116    self.spans.iter().filter(move |s| s.name == name)
117  }
118}
119
120/// A trace being built.
121struct Open {
122  id: u64,
123  started: Instant,
124  spans: Vec<Span>,
125  /// Span id to its index in `spans`, so a child finds its parent.
126  index: HashMap<u64, usize>,
127  /// Every span id the trace took, so the root's close can clear them all
128  /// out of `belongs` rather than leaking one per span that outlives it.
129  ids: Vec<u64>,
130  /// How many spans are still open, so the root closing is unambiguous.
131  root: u64,
132}
133
134/// The handle a caller reads through, and what the layer writes into.
135#[derive(Clone)]
136pub struct Traces {
137  inner: Arc<Inner>,
138}
139
140struct Inner {
141  next_id: AtomicU64,
142  /// Traces still running, by the id of their root span.
143  open: Mutex<HashMap<u64, Open>>,
144  /// Every open span's root, so a deep span finds its trace in one lookup.
145  belongs: Mutex<HashMap<u64, u64>>,
146  finished: Mutex<Vec<Trace>>,
147  ring: usize,
148  listeners: Mutex<Vec<Arc<dyn Fn(&Trace) + Send + Sync>>>,
149}
150
151impl Traces {
152  /// The trace the current span belongs to, as far as it has got. `None`
153  /// outside a root.
154  pub fn current(&self) -> Option<Trace> {
155    let id = tracing::Span::current().id()?;
156    self.of_span(id.into_u64())
157  }
158
159  /// The trace a span id belongs to, for a caller holding one.
160  pub fn of_span(&self, span: u64) -> Option<Trace> {
161    let root = *self.inner.belongs.lock().get(&span)?;
162    let open = self.inner.open.lock();
163    let held = open.get(&root)?;
164    Some(Trace { id: held.id, spans: held.spans.clone(), duration: Duration::ZERO })
165  }
166
167  /// The most recently finished traces, newest last, at most `n`.
168  pub fn recent(&self, n: usize) -> Vec<Trace> {
169    let held = self.inner.finished.lock();
170    let from = held.len().saturating_sub(n);
171    held[from..].to_vec()
172  }
173
174  /// How many traces the ring is holding.
175  pub fn len(&self) -> usize {
176    self.inner.finished.lock().len()
177  }
178
179  pub fn is_empty(&self) -> bool {
180    self.len() == 0
181  }
182
183  /// Drops every finished trace. What a test calls between cases.
184  pub fn clear(&self) {
185    self.inner.finished.lock().clear();
186  }
187
188  /// Called with each trace as its root closes, before it enters the ring.
189  /// This is where an exporter goes, and where a caller that wants to keep
190  /// only the slow or failed ones decides, which nothing sampling at span
191  /// start can do.
192  pub fn on_finish(&self, f: impl Fn(&Trace) + Send + Sync + 'static) {
193    self.inner.listeners.lock().push(Arc::new(f));
194  }
195}
196
197/// The layer and the handle to read it through. Compose the layer into a
198/// registry beside whatever else is listening.
199pub fn layer() -> (TraceLayer, Traces) {
200  with_ring(DEFAULT_RING)
201}
202
203/// `layer` with the ring sized. Zero keeps nothing, which is what a process
204/// that only exports wants.
205pub fn with_ring(ring: usize) -> (TraceLayer, Traces) {
206  let traces = Traces {
207    inner: Arc::new(Inner {
208      next_id: AtomicU64::new(1),
209      open: Mutex::new(HashMap::new()),
210      belongs: Mutex::new(HashMap::new()),
211      finished: Mutex::new(Vec::new()),
212      ring,
213      listeners: Mutex::new(Vec::new()),
214    }),
215  };
216  (TraceLayer { traces: traces.clone() }, traces)
217}
218
219pub struct TraceLayer {
220  traces: Traces,
221}
222
223impl<S> tracing_subscriber::Layer<S> for TraceLayer
224where
225  S: Subscriber + for<'a> LookupSpan<'a>,
226{
227  /// Never clamp what reaches the registry: a collector wants every span a
228  /// caller opened, whatever the log appenders are configured to write.
229  fn max_level_hint(&self) -> Option<tracing::level_filters::LevelFilter> {
230    Some(tracing::level_filters::LevelFilter::TRACE)
231  }
232
233  fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
234    let mut fields = Fields::default();
235    attrs.record(&mut fields);
236    let is_root = matches!(fields.map.get(ROOT_FIELD), Some(Value::Bool(true)));
237
238    let this = id.into_u64();
239    let parent_span = ctx.span(id).and_then(|s| s.parent().map(|p| p.id().into_u64()));
240    let root = match is_root {
241      true => this,
242      false => match parent_span.and_then(|p| self.traces.inner.belongs.lock().get(&p).copied()) {
243        Some(root) => root,
244        // Not under any root: nothing to collect it into.
245        None => return,
246      },
247    };
248
249    let mut open = self.traces.inner.open.lock();
250    if is_root {
251      let trace_id = self.traces.inner.next_id.fetch_add(1, Ordering::Relaxed);
252      open.insert(this, Open { id: trace_id, started: Instant::now(), spans: Vec::new(), index: HashMap::new(), ids: Vec::new(), root: this });
253    }
254    let Some(held) = open.get_mut(&root) else { return };
255    let parent = parent_span.and_then(|p| held.index.get(&p).copied());
256    let depth = parent.map(|i| held.spans[i].depth + 1).unwrap_or(0);
257    let metadata = attrs.metadata();
258    held.index.insert(this, held.spans.len());
259    held.ids.push(this);
260    held.spans.push(Span {
261      name: metadata.name(),
262      target: metadata.target().to_owned(),
263      level: *metadata.level(),
264      parent,
265      depth,
266      at: held.started.elapsed(),
267      duration: Duration::ZERO,
268      fields: fields.map,
269    });
270    drop(open);
271    self.traces.inner.belongs.lock().insert(this, root);
272  }
273
274  fn on_record(&self, id: &Id, values: &Record<'_>, _ctx: Context<'_, S>) {
275    let this = id.into_u64();
276    let Some(root) = self.traces.inner.belongs.lock().get(&this).copied() else { return };
277    let mut fields = Fields::default();
278    values.record(&mut fields);
279    let mut open = self.traces.inner.open.lock();
280    let Some(held) = open.get_mut(&root) else { return };
281    let Some(at) = held.index.get(&this).copied() else { return };
282    held.spans[at].fields.extend(fields.map);
283  }
284
285  fn on_close(&self, id: Id, _ctx: Context<'_, S>) {
286    let this = id.into_u64();
287    let Some(root) = self.traces.inner.belongs.lock().remove(&this) else { return };
288    let mut open = self.traces.inner.open.lock();
289    let Some(held) = open.get_mut(&root) else { return };
290    if let Some(at) = held.index.get(&this).copied() {
291      held.spans[at].duration = held.started.elapsed().saturating_sub(held.spans[at].at);
292    }
293    if held.root != this {
294      return;
295    }
296    let Some(held) = open.remove(&root) else { return };
297    drop(open);
298    {
299      let mut belongs = self.traces.inner.belongs.lock();
300      for id in &held.ids {
301        belongs.remove(id);
302      }
303    }
304    let trace = Trace { id: held.id, duration: held.started.elapsed(), spans: held.spans };
305    for listener in self.traces.inner.listeners.lock().iter() {
306      listener(&trace);
307    }
308    if self.traces.inner.ring == 0 {
309      return;
310    }
311    let mut finished = self.traces.inner.finished.lock();
312    finished.push(trace);
313    let ring = self.traces.inner.ring;
314    if finished.len() > ring {
315      let over = finished.len() - ring;
316      finished.drain(..over);
317    }
318  }
319}
320
321#[derive(Default)]
322struct Fields {
323  map: HashMap<String, Value>,
324}
325
326impl Visit for Fields {
327  fn record_str(&mut self, field: &Field, value: &str) {
328    self.map.insert(field.name().to_owned(), Value::Str(value.to_owned()));
329  }
330  fn record_i64(&mut self, field: &Field, value: i64) {
331    self.map.insert(field.name().to_owned(), Value::Int(value));
332  }
333  fn record_u64(&mut self, field: &Field, value: u64) {
334    self.map.insert(field.name().to_owned(), Value::UInt(value));
335  }
336  fn record_bool(&mut self, field: &Field, value: bool) {
337    self.map.insert(field.name().to_owned(), Value::Bool(value));
338  }
339  fn record_f64(&mut self, field: &Field, value: f64) {
340    self.map.insert(field.name().to_owned(), Value::Float(value));
341  }
342  fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
343    self.map.insert(field.name().to_owned(), Value::Str(format!("{value:?}")));
344  }
345}