await_tree/
span.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15type SpanName = flexstr::SharedStr;
16
17#[doc(hidden)]
18pub fn fmt_span(args: std::fmt::Arguments<'_>) -> Span {
19    let name = if let Some(str) = args.as_str() {
20        SpanName::from_ref(str)
21    } else {
22        flexstr::flex_fmt(args)
23    };
24    Span::new(name)
25}
26
27/// Creates a new span with formatted name.
28///
29/// [`instrument_await`] accepts any type that implements [`AsRef<str>`] as the span name.
30/// This macro provides similar functionality to [`format!`], but with improved performance
31/// by creating the span name on the stack when possible, avoiding unnecessary allocations.
32///
33/// [`instrument_await`]: crate::InstrumentAwait::instrument_await
34#[macro_export]
35// XXX: Without this extra binding (`let res = ..`), it will make the future `!Send`.
36//      This is also how `std::format!` behaves. But why?
37macro_rules! span {
38    ($($fmt_arg:tt)*) => {{
39        let res = $crate::__private::fmt_span(format_args!($($fmt_arg)*));
40        res
41    }};
42}
43
44/// A cheaply cloneable span in the await-tree.
45#[cfg_attr(feature = "serde", derive(serde::Serialize))]
46#[derive(Debug, Clone, Default, PartialEq, Eq, PartialOrd, Ord)]
47pub struct Span {
48    pub(crate) name: SpanName,
49    pub(crate) is_verbose: bool,
50    pub(crate) is_long_running: bool,
51}
52
53impl Span {
54    fn new(name: SpanName) -> Self {
55        Self {
56            name,
57            is_verbose: false,
58            is_long_running: false,
59        }
60    }
61}
62
63impl Span {
64    /// Set the verbose attribute of the span.
65    ///
66    /// When a span is marked as verbose, it will be included in the output
67    /// only if the `verbose` flag in the [`Config`] is set.
68    ///
69    /// [`Config`]: crate::Config
70    pub fn verbose(mut self) -> Self {
71        self.is_verbose = true;
72        self
73    }
74
75    /// Set the long-running attribute of the span.
76    ///
77    /// When a span is marked as long-running, it will not be marked as "!!!"
78    /// in the formatted [`Tree`] if it takes too long to complete. The root
79    /// span is always marked as long-running.
80    ///
81    /// [`Tree`]: crate::Tree
82    pub fn long_running(mut self) -> Self {
83        self.is_long_running = true;
84        self
85    }
86}
87
88/// Convert a value into a span and set attributes.
89#[easy_ext::ext(SpanExt)]
90impl<T: Into<Span>> T {
91    /// Convert `self` into a span and set the verbose attribute.
92    ///
93    /// See [`Span::verbose`] for more details.
94    pub fn verbose(self) -> Span {
95        self.into().verbose()
96    }
97
98    /// Convert `self` into a span and set the long-running attribute.
99    ///
100    /// See [`Span::long_running`] for more details.
101    pub fn long_running(self) -> Span {
102        self.into().long_running()
103    }
104}
105
106impl<S: AsRef<str>> From<S> for Span {
107    fn from(value: S) -> Self {
108        Self::new(SpanName::from_ref(value))
109    }
110}
111
112impl std::fmt::Display for Span {
113    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114        self.name.fmt(f)
115    }
116}
await_tree/span.rs

await_tree/
span.rs
