apalis_core/backend/ext/
instrument.rs

1use futures_sink::Sink;
2use futures_util::SinkExt;
3use std::pin::Pin;
4use std::task::{Context, Poll};
5
6use tracing::Span;
7
8use crate::{backend::*, task::Task, worker::context::WorkerContext};
9
10/// Instruments the inner [Backend] with the provided `Span`, returning an
11/// `Instrumented` wrapper.
12#[derive(Debug, Clone)]
13pub struct Instrumented<B> {
14    backend: B,
15    span: Span,
16}
17
18impl<B> Instrumented<B> {
19    /// Create a new `Instrumented` backend with the given tracing span.
20    pub fn new(backend: B, span: Span) -> Self {
21        Self { backend, span }
22    }
23
24    /// Returns the wrapped backend.
25    pub fn inner(&self) -> &B {
26        &self.backend
27    }
28
29    /// Returns a mutable reference to the wrapped backend.
30    pub fn inner_mut(&mut self) -> &mut B {
31        &mut self.backend
32    }
33
34    /// Returns the tracing span.
35    pub fn span(&self) -> &Span {
36        &self.span
37    }
38
39    /// Unwraps the instrumented backend.
40    pub fn into_inner(self) -> B {
41        self.backend
42    }
43}
44
45impl<B> Backend for Instrumented<B>
46where
47    B: Backend,
48{
49    type Task = B::Task;
50    type Error = B::Error;
51
52    fn poll_ready(
53        &mut self,
54        cx: &mut Context<'_>,
55        worker: &WorkerContext,
56    ) -> Poll<Result<(), Self::Error>> {
57        let _entered = self.span.enter();
58        self.backend.poll_ready(cx, worker)
59    }
60
61    fn poll_next(
62        &mut self,
63        cx: &mut Context<'_>,
64        worker: &WorkerContext,
65    ) -> Poll<Option<Result<Self::Task, Self::Error>>> {
66        let _entered = self.span.enter();
67        self.backend.poll_next(cx, worker)
68    }
69
70    fn poll_close(
71        &mut self,
72        cx: &mut Context<'_>,
73        worker: &WorkerContext,
74    ) -> Poll<Result<(), Self::Error>> {
75        let _entered = self.span.enter();
76        self.backend.poll_close(cx, worker)
77    }
78}
79
80impl<B, T, Err> Sink<T> for Instrumented<B>
81where
82    B: Sink<T, Error = Err> + Unpin,
83{
84    type Error = Err;
85    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
86        let this = self.get_mut();
87        let _entered = this.span.enter();
88        this.backend.start_send_unpin(item)
89    }
90    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
91        let this = self.get_mut();
92        let _entered = this.span.enter();
93        this.backend.poll_ready_unpin(cx)
94    }
95    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
96        let this = self.get_mut();
97        let _entered = this.span.enter();
98        this.backend.poll_flush_unpin(cx)
99    }
100    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
101        let this = self.get_mut();
102        let _entered = this.span.enter();
103        this.backend.poll_close_unpin(cx)
104    }
105}
106
107delegate_expose!(
108    impl<B> for Instrumented<B>
109    where {
110        B: Send + Sync,
111        B::Task: Send + Clone + 'static,
112    }
113    => backend
114);
115
116delegate_config!(Instrumented<B>, backend);
117delegate_deref!(Instrumented<B>, backend);
118delegate_codec!(Instrumented<B>, backend);
apalis_core/backend/ext/instrument.rs

apalis_core/backend/ext/
instrument.rs
