Skip to main content

tdm_server_rust/telemetry/skywalking/
context.rs

1//! 请求级 SkyWalking TracingContext(tokio task_local)
2
3use ::skywalking::trace::{
4    span::{HandleSpanObject, Span},
5    trace_context::TracingContext,
6    tracer::create_trace_context,
7};
8use std::cell::RefCell;
9use std::future::Future;
10
11tokio::task_local! {
12    static SW_TRACE_CTX: RefCell<Option<TracingContext>>;
13}
14
15/// 读取当前 SW trace_id(32hex 字符串)
16pub fn current_trace_id() -> Option<String> {
17    SW_TRACE_CTX
18        .try_with(|cell: &RefCell<Option<TracingContext>>| {
19            cell.borrow()
20                .as_ref()
21                .map(|ctx: &TracingContext| ctx.trace_id().to_string())
22        })
23        .ok()
24        .flatten()
25}
26
27/// 是否处于 HTTP 请求 SW 上下文内
28pub fn is_in_request_context() -> bool {
29    SW_TRACE_CTX
30        .try_with(|cell: &RefCell<Option<TracingContext>>| cell.borrow().is_some())
31        .unwrap_or(false)
32}
33
34/// 在 SW 上下文作用域内执行 future;结束时 wait 全部 async span 再释放 ctx
35pub async fn with_context<F, R>(ctx: TracingContext, fut: F) -> R
36where
37    F: Future<Output = R>,
38{
39    SW_TRACE_CTX
40        .scope(RefCell::new(Some(ctx)), async move {
41            let result = fut.await;
42            if let Some(c) = SW_TRACE_CTX
43                .try_with(|cell: &RefCell<Option<TracingContext>>| cell.borrow_mut().take())
44                .ok()
45                .flatten()
46            {
47                c.wait();
48            }
49            result
50        })
51        .await
52}
53
54/// 可变访问当前 TracingContext
55pub fn with_current_mut<F, R>(f: F) -> Option<R>
56where
57    F: FnOnce(&mut TracingContext) -> R,
58{
59    SW_TRACE_CTX
60        .try_with(|cell: &RefCell<Option<TracingContext>>| {
61            cell.borrow_mut().as_mut().map(f)
62        })
63        .ok()
64        .flatten()
65}
66
67/// 只读访问当前 TracingContext
68pub fn with_current<F, R>(f: F) -> Option<R>
69where
70    F: FnOnce(&TracingContext) -> R,
71{
72    SW_TRACE_CTX
73        .try_with(|cell: &RefCell<Option<TracingContext>>| cell.borrow().as_ref().map(f))
74        .ok()
75        .flatten()
76}
77
78/// 创建 LocalSpan(需在 Entry 上下文内;无 Entry 时返回 None 而非 panic)
79pub fn create_local_span(operation: &str) -> Option<Span> {
80    std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
81        with_current_mut(|ctx| ctx.create_local_span(operation))
82    }))
83    .ok()
84    .flatten()
85}
86
87/// 后台任务独立 TraceContext(tag layer=bootstrap,不进 HTTP 日志)
88pub async fn with_bootstrap_context<F, R>(operation: &str, tag: &str, fut: F) -> R
89where
90    F: Future<Output = R>,
91{
92    if !super::reporter::is_ready() {
93        return fut.await;
94    }
95    let mut ctx = create_trace_context();
96    let mut entry = ctx.create_entry_span(operation);
97    entry.add_tag("layer", "bootstrap");
98    entry.add_tag("task", tag);
99    with_context(ctx, async move {
100        let _entry = entry;
101        fut.await
102    })
103    .await
104}