tdm_server_rust/telemetry/skywalking/
context.rs1use ::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
15pub 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
27pub 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
34pub 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
54pub 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
67pub 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
78pub 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
87pub 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}