await_tree/
lib.rs

1// Copyright 2023 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
15//! Instrument await-tree for actor-based applications.
16
17#![forbid(missing_docs)]
18
19use std::future::Future;
20
21mod context;
22mod future;
23mod obj_utils;
24mod registry;
25mod root;
26mod spawn;
27
28pub use context::*;
29pub use future::*;
30pub use registry::*;
31pub use root::*;
32pub use spawn::*;
33
34/// A cheaply cloneable span in the await-tree.
35#[derive(Debug, Clone, Default, PartialEq, Eq, PartialOrd, Ord)]
36pub struct Span(flexstr::SharedStr);
37
38impl Span {
39    pub(crate) fn as_str(&self) -> &str {
40        self.0.as_str()
41    }
42}
43
44impl<S: AsRef<str>> From<S> for Span {
45    fn from(value: S) -> Self {
46        Self(flexstr::SharedStr::from_ref(value))
47    }
48}
49
50impl std::fmt::Display for Span {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        self.0.fmt(f)
53    }
54}
55
56/// Attach spans to a future to be traced in the await-tree.
57pub trait InstrumentAwait: Future + Sized {
58    /// Instrument the future with a span.
59    fn instrument_await(self, span: impl Into<Span>) -> Instrumented<Self, false> {
60        Instrumented::new(self, span.into())
61    }
62
63    /// Instrument the future with a verbose span, which is optionally enabled based on the registry
64    /// configuration.
65    fn verbose_instrument_await(self, span: impl Into<Span>) -> Instrumented<Self, true> {
66        Instrumented::new(self, span.into())
67    }
68}
69impl<F> InstrumentAwait for F where F: Future {}
70
71#[cfg(test)]
72mod tests;
await_tree/lib.rs

await_tree/
lib.rs
