apalis_core/backend/ext/lifecycle/
after_start.rs

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

apalis_core/backend/ext/lifecycle/
after_start.rs
