log_args_runtime/
lib.rs

1use std::cell::RefCell;
2use std::collections::HashMap;
3use std::future::Future;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6
7pub const WITH_CONTEXT_ENABLED: bool = cfg!(feature = "with_context");
8
9thread_local! {
10    static CONTEXT_STACK: RefCell<Vec<HashMap<String, String>>> = const { RefCell::new(Vec::new()) };
11    static ASYNC_STACK_TL: RefCell<Vec<HashMap<String, String>>> = const { RefCell::new(Vec::new()) };
12}
13
14pub fn set_global_context(_: &str, _: &str) {}
15pub fn get_global_context() -> Option<HashMap<String, String>> {
16    None
17}
18
19pub fn get_context_merged() -> HashMap<String, String> {
20    let mut ctx = ASYNC_STACK_TL.with(|s| {
21        let b = s.borrow();
22        b.iter().fold(HashMap::new(), |mut acc, c| {
23            acc.extend(c.clone());
24            acc
25        })
26    });
27    ctx.extend(CONTEXT_STACK.with(|s| {
28        let b = s.borrow();
29        b.iter().fold(HashMap::new(), |mut acc, c| {
30            acc.extend(c.clone());
31            acc
32        })
33    }));
34    ctx
35}
36
37pub fn get_context() -> HashMap<String, String> {
38    get_context_merged()
39}
40pub fn get_async_context() -> HashMap<String, String> {
41    get_context_merged()
42}
43pub fn get_context_all() -> HashMap<String, String> {
44    get_context_merged()
45}
46pub fn get_context_value(key: &str) -> Option<String> {
47    get_context_merged().get(key).cloned()
48}
49pub fn get_context_value_merged(key: &str) -> Option<String> {
50    get_context_value(key)
51}
52
53pub struct ContextFut<F> {
54    fut: F,
55    stack: Vec<HashMap<String, String>>,
56}
57
58impl<F: Future> Future for ContextFut<F> {
59    type Output = F::Output;
60    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
61        let this = unsafe { self.get_unchecked_mut() };
62        let old_stack = ASYNC_STACK_TL.with(|s| {
63            let old = s.borrow().clone();
64            *s.borrow_mut() = this.stack.clone();
65            old
66        });
67        let res = unsafe { Pin::new_unchecked(&mut this.fut) }.poll(cx);
68        ASYNC_STACK_TL.with(|s| *s.borrow_mut() = old_stack);
69        res
70    }
71}
72
73pub async fn with_async_context<F, R>(new_context: HashMap<String, String>, fut: F) -> R
74where
75    F: Future<Output = R>,
76{
77    let mut stack = ASYNC_STACK_TL.with(|s| s.borrow().clone());
78    stack.push(new_context);
79    ContextFut { fut, stack }.await
80}
81
82pub fn scope_current_context<F: Future>(fut: F) -> impl Future<Output = F::Output> {
83    let stack = ASYNC_STACK_TL.with(|s| s.borrow().clone());
84    ContextFut { fut, stack }
85}
86
87pub fn instrument_spawn<F: Future>(fut: F) -> impl Future<Output = F::Output> {
88    scope_current_context(::tracing::Instrument::instrument(
89        fut,
90        ::tracing::Span::current(),
91    ))
92}
93
94#[macro_export]
95macro_rules! log_with_context {
96    ($log_macro:path, $context:expr, $($args:tt)*) => {
97        {
98            let ctx = $context;
99            if !$crate::WITH_CONTEXT_ENABLED {
100                $log_macro!($($args)*);
101            } else {
102                $log_macro!(context = ?ctx, $($args)*);
103            }
104        }
105    };
106}
107
108#[macro_export]
109macro_rules! info { ($($t:tt)*) => { $crate::log_with_context!(::tracing::info, $crate::get_context_merged(), $($t)*) }; }
110#[macro_export]
111macro_rules! warn { ($($t:tt)*) => { $crate::log_with_context!(::tracing::warn, $crate::get_context_merged(), $($t)*) }; }
112#[macro_export]
113macro_rules! error { ($($t:tt)*) => { $crate::log_with_context!(::tracing::error, $crate::get_context_merged(), $($t)*) }; }
114#[macro_export]
115macro_rules! debug { ($($t:tt)*) => { $crate::log_with_context!(::tracing::debug, $crate::get_context_merged(), $($t)*) }; }
116#[macro_export]
117macro_rules! trace { ($($t:tt)*) => { $crate::log_with_context!(::tracing::trace, $crate::get_context_merged(), $($t)*) }; }
118
119pub struct ContextGuard;
120impl Drop for ContextGuard {
121    fn drop(&mut self) {
122        CONTEXT_STACK.with(|s| {
123            s.borrow_mut().pop();
124        });
125    }
126}
127pub fn push_context(c: HashMap<String, String>) -> ContextGuard {
128    CONTEXT_STACK.with(|s| {
129        s.borrow_mut().push(c);
130    });
131    ContextGuard
132}
133pub fn capture_context() -> ContextGuard {
134    push_context(get_context_merged())
135}
136
137pub fn get_inherited_context_string() -> String {
138    let ctx = get_context_merged();
139    ctx.iter()
140        .map(|(k, v)| format!("{}={}", k, v))
141        .collect::<Vec<_>>()
142        .join(",")
143}
144
145pub fn get_inherited_fields_map() -> HashMap<String, String> {
146    get_context_merged()
147}
148
149#[doc(hidden)]
150pub struct AsyncContextGuard;
151impl Drop for AsyncContextGuard {
152    fn drop(&mut self) {}
153}
154#[doc(hidden)]
155pub fn push_async_context(_: HashMap<String, String>) -> AsyncContextGuard {
156    AsyncContextGuard
157}
158#[doc(hidden)]
159pub fn get_current_async_stack() -> Vec<HashMap<String, String>> {
160    ASYNC_STACK_TL.with(|s| s.borrow().clone())
161}
162#[doc(hidden)]
163pub fn get_current_task_context() -> HashMap<String, String> {
164    get_context_merged()
165}
166#[doc(hidden)]
167pub static TASK_CONTEXT_STACK: () = ();
168#[doc(hidden)]
169pub static TASK_CONTEXT_STACK_UNIT: () = ();
log_args_runtime/lib.rs

log_args_runtime/
lib.rs
