apalis_core/backend/ext/lifecycle/
after_stop.rs

1use std::{
2    future::Future,
3    ops::{Deref, DerefMut},
4    pin::Pin,
5    sync::Arc,
6    task::{
7        Context,
8        Poll::{self},
9    },
10};
11
12use futures_sink::Sink;
13use futures_util::{FutureExt, SinkExt};
14
15use crate::{
16    backend::{
17        ext::lifecycle::{FnHandler, HookState},
18        future::BoxSyncFuture,
19        *,
20    },
21    task::Task,
22    worker::context::WorkerContext,
23};
24
25/// A hook that runs after a backend has stopped.
26///
27/// `AfterStop` allows an asynchronous callback to perform cleanup or other
28/// work on the backend after it has stopped.
29///
30/// The callback receives a mutable reference to the backend and must return a
31/// future resolving to `Result<(), B::Error>`.
32pub struct AfterStop<B, FutErr> {
33    pub(super) backend: B,
34    pub(super) f: FnHandler<B, FutErr>,
35    state: HookState<FutErr>,
36}
37
38impl<B, FutErr> std::fmt::Debug for AfterStop<B, FutErr>
39where
40    B: std::fmt::Debug,
41{
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        f.debug_struct("AfterStop")
44            .field("backend", &self.backend)
45            .field("f", &"<function>")
46            .field("state", &self.state)
47            .finish()
48    }
49}
50
51impl<B, FutErr> Clone for AfterStop<B, FutErr>
52where
53    B: Clone,
54{
55    fn clone(&self) -> Self {
56        Self {
57            backend: self.backend.clone(),
58            f: self.f.clone(),
59            state: HookState::Pending,
60        }
61    }
62}
63
64impl<B: Backend> AfterStop<B, B::Error> {
65    /// Creates a new `AfterStop` hook.
66    ///
67    /// The provided function is called with a mutable reference to the
68    /// backend after it has stopped.
69    ///
70    /// # Errors
71    ///
72    /// If the callback returns an error, the error is propagated to the
73    /// backend lifecycle.
74    pub fn new<F, Fut>(backend: B, f: F) -> Self
75    where
76        F: Fn(&mut B) -> Fut + 'static + Send + Sync,
77        Fut: Future<Output = Result<(), B::Error>> + Send + 'static,
78    {
79        Self {
80            backend,
81            f: Arc::new(move |b: &mut B| BoxSyncFuture::new(f(b).boxed())),
82            state: HookState::Pending,
83        }
84    }
85}
86
87impl<B> Backend for AfterStop<B, B::Error>
88where
89    B: Backend,
90{
91    type Task = B::Task;
92    type Error = B::Error;
93
94    fn poll_ready(
95        &mut self,
96        cx: &mut Context<'_>,
97        worker: &WorkerContext,
98    ) -> Poll<Result<(), Self::Error>> {
99        self.backend.poll_ready(cx, worker)
100    }
101
102    fn poll_next(
103        &mut self,
104        cx: &mut Context<'_>,
105        worker: &WorkerContext,
106    ) -> Poll<Option<Result<Self::Task, Self::Error>>> {
107        self.backend.poll_next(cx, worker)
108    }
109
110    fn poll_close(
111        &mut self,
112        cx: &mut Context<'_>,
113        worker: &WorkerContext,
114    ) -> Poll<Result<(), Self::Error>> {
115        match self.backend.poll_close(cx, worker) {
116            Poll::Ready(Ok(())) => {}
117            Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
118            Poll::Pending => return Poll::Pending,
119        }
120        let backend = &mut self.backend;
121        let f = &self.f;
122        self.state.poll_hook(cx, || f(backend))
123    }
124}
125
126impl<B, FutErr> Deref for AfterStop<B, FutErr> {
127    type Target = B;
128
129    fn deref(&self) -> &Self::Target {
130        &self.backend
131    }
132}
133
134impl<B, FutErr> DerefMut for AfterStop<B, FutErr> {
135    fn deref_mut(&mut self) -> &mut Self::Target {
136        &mut self.backend
137    }
138}
139
140delegate_sink!(AfterStop<B, FutErr>, backend);
141
142delegate_codec!(AfterStop<B, FutErr>, backend);
143
144delegate_config!(AfterStop<B, FutErr>, backend);
145
146delegate_expose!(
147    impl<B> for AfterStop<B, B::Error>
148    where {
149        B: Backend + Send,
150
151    }
152    => backend
153);
apalis_core/backend/ext/lifecycle/after_stop.rs

apalis_core/backend/ext/lifecycle/
after_stop.rs
