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#[derive(Debug, Clone, Default, PartialEq, Eq, PartialOrd, Ord)]
46pub struct Span {
47    pub(crate) name: SpanName,
48    pub(crate) is_verbose: bool,
49    pub(crate) is_long_running: bool,
50}
51
52impl Span {
53    fn new(name: SpanName) -> Self {
54        Self {
55            name,
56            is_verbose: false,
57            is_long_running: false,
58        }
59    }
60}
61
62impl Span {
63    /// Set the verbose attribute of the span.
64    ///
65    /// When a span is marked as verbose, it will be included in the output
66    /// only if the `verbose` flag in the [`Config`] is set.
67    ///
68    /// [`Config`]: crate::Config
69    pub fn verbose(mut self) -> Self {
70        self.is_verbose = true;
71        self
72    }
73
74    /// Set the long-running attribute of the span.
75    ///
76    /// When a span is marked as long-running, it will not be marked as "!!!"
77    /// in the formatted [`Tree`] if it takes too long to complete. The root
78    /// span is always marked as long-running.
79    ///
80    /// [`Tree`]: crate::Tree
81    pub fn long_running(mut self) -> Self {
82        self.is_long_running = true;
83        self
84    }
85}
86
87/// Convert a value into a span and set attributes.
88#[easy_ext::ext(SpanExt)]
89impl<T: Into<Span>> T {
90    /// Convert `self` into a span and set the verbose attribute.
91    ///
92    /// See [`Span::verbose`] for more details.
93    pub fn verbose(self) -> Span {
94        self.into().verbose()
95    }
96
97    /// Convert `self` into a span and set the long-running attribute.
98    ///
99    /// See [`Span::long_running`] for more details.
100    pub fn long_running(self) -> Span {
101        self.into().long_running()
102    }
103}
104
105impl<S: AsRef<str>> From<S> for Span {
106    fn from(value: S) -> Self {
107        Self::new(SpanName::from_ref(value))
108    }
109}
110
111impl std::fmt::Display for Span {
112    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
113        self.name.fmt(f)
114    }
115}
await_tree/span.rs

await_tree/
span.rs
