Skip to main content

teloxide_ng/dispatching/
tracing.rs

1use super::UpdateHandler;
2
3use dptree::{
4    HandlerSignature,
5    di::{Asyncify, Injectable},
6    prelude::DependencyMap,
7};
8use std::{collections::BTreeSet, sync::Arc};
9use tracing::{Instrument, Span};
10
11pub trait UpdateHandlerTracingExt<E> {
12    /// Returns an `UpdateHandler` wrapped in an async span.
13    fn instrument_with_async<F, FnArgs>(self, f: F) -> Self
14    where
15        F: Injectable<Span, FnArgs> + Send + Sync + 'static;
16
17    /// Returns an `UpdateHandler` wrapped in a span.
18    fn instrument_with<F, FnArgs>(self, f: F) -> Self
19    where
20        Asyncify<F>: Injectable<Span, FnArgs> + Send + Sync + 'static;
21}
22
23impl<E: 'static> UpdateHandlerTracingExt<E> for UpdateHandler<E> {
24    fn instrument_with_async<F, FnArgs>(self, f: F) -> UpdateHandler<E>
25    where
26        F: Injectable<Span, FnArgs> + Send + Sync + 'static,
27    {
28        let f = Arc::new(f);
29
30        dptree::from_fn_with_description(
31            self.description().clone(),
32            move |deps: DependencyMap, cont| {
33                let self_c = self.clone();
34                let f = Arc::clone(&f);
35
36                async move {
37                    let f = f.inject(&deps);
38                    let span = f().await;
39                    drop(f);
40
41                    self_c.execute(deps, cont).instrument(span).await
42                }
43            },
44            HandlerSignature::Other {
45                obligations: F::obligations(),
46                guaranteed_outcomes: BTreeSet::new(),
47                conditional_outcomes: BTreeSet::new(),
48                continues: true,
49            },
50        )
51    }
52
53    fn instrument_with<F, FnArgs>(self, f: F) -> Self
54    where
55        Asyncify<F>: Injectable<Span, FnArgs> + Send + Sync + 'static,
56    {
57        self.instrument_with_async(Asyncify(f))
58    }
59}