log_io/
context.rs

1//! Thread-local logging context.
2//!
3//! Context carries ambient key-value data that gets attached to every
4//! record emitted from the current thread. Typical use: stash a
5//! `trace_id` / `request_id` at the entry of a request handler so all
6//! downstream log lines correlate. The [`ContextGuard`] returned by
7//! [`with_field`] (and friends) removes the value when dropped.
8//!
9//! Context is intentionally thread-scoped. For async task propagation,
10//! capture context fields at task spawn and re-install them on the
11//! task's executor thread; this crate does not depend on a specific
12//! async runtime.
13
14use std::cell::RefCell;
15use std::sync::atomic::{AtomicU64, Ordering};
16
17use crate::record::Field;
18use crate::value::Value;
19
20thread_local! {
21    static CONTEXT: RefCell<ContextSlots> = const { RefCell::new(ContextSlots::new()) };
22}
23
24/// Maximum number of context slots per thread. Fixed-size so the
25/// snapshot path is allocation-free.
26pub const MAX_CONTEXT_SLOTS: usize = 16;
27
28#[derive(Debug)]
29struct ContextSlot {
30    key: String,
31    value: ContextValue,
32    id: u64,
33}
34
35#[derive(Debug)]
36enum ContextValue {
37    Null,
38    Bool(bool),
39    I64(i64),
40    U64(u64),
41    F64(f64),
42    Str(String),
43    Char(char),
44}
45
46impl ContextValue {
47    fn as_value(&self) -> Value<'_> {
48        match self {
49            Self::Null => Value::Null,
50            Self::Bool(b) => Value::Bool(*b),
51            Self::I64(n) => Value::I64(*n),
52            Self::U64(n) => Value::U64(*n),
53            Self::F64(n) => Value::F64(*n),
54            Self::Str(s) => Value::Str(s.as_str()),
55            Self::Char(c) => Value::Char(*c),
56        }
57    }
58}
59
60fn owned_from_value(v: Value<'_>) -> ContextValue {
61    match v {
62        Value::Null => ContextValue::Null,
63        Value::Bool(b) => ContextValue::Bool(b),
64        Value::I64(n) => ContextValue::I64(n),
65        Value::U64(n) => ContextValue::U64(n),
66        Value::F64(n) => ContextValue::F64(n),
67        Value::Str(s) => ContextValue::Str(s.to_owned()),
68        Value::Char(c) => ContextValue::Char(c),
69    }
70}
71
72#[derive(Debug)]
73struct ContextSlots {
74    slots: Vec<ContextSlot>,
75}
76
77impl ContextSlots {
78    const fn new() -> Self {
79        Self { slots: Vec::new() }
80    }
81}
82
83static NEXT_ID: AtomicU64 = AtomicU64::new(1);
84
85/// RAII guard that removes a context slot when dropped.
86///
87/// Returned by [`with_field`], [`with_trace_id`], and
88/// [`with_request_id`]. Drop the guard at the end of the scope where
89/// the context applies. Dropping guards out-of-order is supported.
90#[must_use = "context is removed when this guard is dropped; bind it to a variable"]
91#[derive(Debug)]
92pub struct ContextGuard {
93    id: u64,
94}
95
96impl Drop for ContextGuard {
97    fn drop(&mut self) {
98        CONTEXT.with(|cell| {
99            if let Ok(mut ctx) = cell.try_borrow_mut() {
100                ctx.slots.retain(|s| s.id != self.id);
101            }
102        });
103    }
104}
105
106/// Push a key/value pair onto the current thread's context.
107///
108/// Returns a guard; when the guard drops, the value is removed. If the
109/// fixed-size slot limit ([`MAX_CONTEXT_SLOTS`]) is exceeded, the
110/// oldest slot is dropped to make room.
111///
112/// # Example
113///
114/// ```
115/// use log_io::context;
116/// use log_io::Value;
117///
118/// {
119///     let _g = context::with_field("user_id", Value::U64(42));
120///     // log calls in this scope carry user_id=42
121/// }
122/// // user_id is gone here.
123/// ```
124pub fn with_field<'a, V: Into<Value<'a>>>(key: &str, value: V) -> ContextGuard {
125    let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
126    let owned = owned_from_value(value.into());
127    CONTEXT.with(|cell| {
128        let mut ctx = cell.borrow_mut();
129        if ctx.slots.len() >= MAX_CONTEXT_SLOTS {
130            ctx.slots.remove(0);
131        }
132        ctx.slots.push(ContextSlot {
133            key: key.to_owned(),
134            value: owned,
135            id,
136        });
137    });
138    ContextGuard { id }
139}
140
141/// Convenience: push a `trace_id` field.
142pub fn with_trace_id(id: &str) -> ContextGuard {
143    with_field("trace_id", Value::Str(id))
144}
145
146/// Convenience: push a `request_id` field.
147pub fn with_request_id(id: &str) -> ContextGuard {
148    with_field("request_id", Value::Str(id))
149}
150
151/// Run `f` with the current thread's context exposed as a borrowed
152/// slice of [`Field`]. No allocation: the slice is filled from a
153/// stack-resident buffer.
154///
155/// Reentrancy-safe: if `f` is invoked transitively from inside another
156/// `with_snapshot` call (or while a `with_field` guard is being
157/// constructed), the inner snapshot sees an empty slice rather than
158/// panicking. This makes the context system safe to use inside sinks
159/// that happen to emit log records of their own.
160pub fn with_snapshot<R>(f: impl FnOnce(&[Field<'_>]) -> R) -> R {
161    CONTEXT.with(|cell| match cell.try_borrow() {
162        Ok(ctx) => {
163            const EMPTY: Field<'static> = Field {
164                key: "",
165                value: Value::Null,
166            };
167            let mut buf: [Field<'_>; MAX_CONTEXT_SLOTS] = [EMPTY; MAX_CONTEXT_SLOTS];
168            let len = ctx.slots.len().min(MAX_CONTEXT_SLOTS);
169            for (i, slot) in ctx.slots.iter().take(len).enumerate() {
170                buf[i] = Field::new(slot.key.as_str(), slot.value.as_value());
171            }
172            f(&buf[..len])
173        }
174        Err(_) => f(&[]),
175    })
176}
177
178/// Remove all context entries on the current thread.
179///
180/// Intended for tests; production code should rely on
181/// [`ContextGuard`] drop semantics.
182pub fn clear() {
183    CONTEXT.with(|cell| {
184        if let Ok(mut ctx) = cell.try_borrow_mut() {
185            ctx.slots.clear();
186        }
187    });
188}
189
190/// Number of context entries on the current thread.
191pub fn len() -> usize {
192    CONTEXT.with(|cell| cell.borrow().slots.len())
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198
199    #[test]
200    fn context_is_isolated_per_thread() {
201        clear();
202        let _g = with_field("k", Value::U64(1));
203        assert_eq!(len(), 1);
204
205        std::thread::spawn(|| {
206            assert_eq!(len(), 0);
207        })
208        .join()
209        .unwrap();
210        clear();
211    }
212
213    #[test]
214    fn guard_removes_on_drop() {
215        clear();
216        {
217            let _g = with_field("k", Value::U64(1));
218            with_snapshot(|f| assert_eq!(f.len(), 1));
219        }
220        with_snapshot(|f| assert_eq!(f.len(), 0));
221    }
222
223    #[test]
224    fn nested_scopes_layer_and_unwind() {
225        clear();
226        let outer = with_field("outer", Value::U64(1));
227        {
228            let _inner = with_field("inner", Value::U64(2));
229            with_snapshot(|f| assert_eq!(f.len(), 2));
230        }
231        with_snapshot(|f| {
232            assert_eq!(f.len(), 1);
233            assert_eq!(f[0].key, "outer");
234        });
235        drop(outer);
236        with_snapshot(|f| assert_eq!(f.len(), 0));
237    }
238
239    #[test]
240    fn slot_overflow_evicts_oldest() {
241        clear();
242        let mut guards = Vec::new();
243        for i in 0..(MAX_CONTEXT_SLOTS + 5) {
244            guards.push(with_field("k", Value::U64(i as u64)));
245        }
246        with_snapshot(|f| {
247            assert_eq!(f.len(), MAX_CONTEXT_SLOTS);
248            if let Value::U64(n) = f[0].value {
249                assert!(n >= 5, "oldest entries should have been evicted");
250            } else {
251                panic!("unexpected variant");
252            }
253        });
254        drop(guards);
255        with_snapshot(|f| assert_eq!(f.len(), 0));
256    }
257
258    #[test]
259    fn trace_and_request_id_helpers() {
260        clear();
261        let _t = with_trace_id("abc");
262        let _r = with_request_id("xyz");
263        with_snapshot(|f| {
264            assert_eq!(f[0].key, "trace_id");
265            assert_eq!(f[0].value, Value::Str("abc"));
266            assert_eq!(f[1].key, "request_id");
267            assert_eq!(f[1].value, Value::Str("xyz"));
268        });
269        clear();
270    }
271}
log_io/context.rs

log_io/
context.rs
