1use std::{
16 task::{Context, Poll},
17 time::Duration,
18};
19
20use futures_core::Stream;
21
22#[cfg(feature = "tracing")]
23use crate::backend::ext::instrument::Instrumented;
24use crate::{
25 backend::{
26 Backend, BackendConfig, WireFormatBackend,
27 codec::Codec,
28 ext::{
29 inspect_err::InspectErr,
30 interleave::Interleave,
31 lifecycle::{AfterStart, AfterStop, BeforeStart, BeforeStop},
32 map_err::MapErr,
33 pipe::Pipe,
34 poll_strategy::{PollStrategy, PollWith, StreamStrategy},
35 wake_on_push::WakeOnPush,
36 with_codec::WithCodec,
37 },
38 },
39 error::BoxDynError,
40 task::Task,
41 worker::context::WorkerContext,
42};
43
44#[cfg(feature = "shared")]
45use crate::backend::ext::shared::Shared;
46
47#[cfg(feature = "sleep")]
48use crate::backend::ext::poll_strategy::{BackoffConfig, BackoffStrategy, IntervalStrategy};
49
50#[macro_use]
51pub mod delegate;
52pub mod inspect_err;
54
55#[cfg(feature = "tracing")]
57pub mod instrument;
58pub mod interleave;
60
61pub mod map_err;
63pub mod pipe;
64pub mod poll_strategy;
65pub mod wake_on_push;
67pub mod with_codec;
68
69pub mod lifecycle;
70
71#[cfg(feature = "shared")]
73pub mod shared;
74
75#[derive(Debug, thiserror::Error)]
77#[non_exhaustive]
78pub enum PollNextArgsError<B: Backend> {
79 #[error("backend error: {0}")]
81 BackendError(B::Error),
82 #[error("failed to decode task args: {0}")]
84 DecodeError(BoxDynError),
85}
86
87pub trait BackendExt: Backend {
89 #[allow(clippy::type_complexity)]
92 fn poll_next_args(
93 &mut self,
94 cx: &mut Context<'_>,
95 worker: &WorkerContext,
96 ) -> Poll<Option<Result<Task<Self::Args>, PollNextArgsError<Self>>>>
97 where
98 Self: Sized + BackendConfig + WireFormatBackend + Backend<Task = Task<Self::Compact>>,
99 Self::Codec: Codec<Self::Args, Compact = Self::Compact>,
100 <Self::Codec as Codec<Self::Args>>::Error: std::error::Error + Send + Sync + 'static,
101 {
102 let next = self.poll_next(cx, worker);
103 let codec = self.codec();
104 next.map(move |item| match item {
105 Some(Ok(task)) => {
106 let task = task.try_map_args(|compact| codec.decode(&compact));
107 Some(task.map_err(|e| PollNextArgsError::DecodeError(e.into())))
108 }
109 Some(Err(e)) => Some(Err(PollNextArgsError::BackendError(e))),
110 None => None,
111 })
112 }
113
114 fn pipe_to<Dst>(self, backend: Dst) -> Pipe<Dst, Self>
120 where
121 Self: Sized,
122 {
123 Pipe::new(self, backend)
124 }
125
126 fn inspect_err<F>(self, f: F) -> InspectErr<Self, F>
128 where
129 Self: Sized,
130 F: Fn(&Self::Error),
131 {
132 InspectErr { backend: self, f }
133 }
134
135 fn map_err<F, E2>(self, f: F) -> MapErr<Self, F>
138 where
139 Self: Sized,
140 F: Fn(Self::Error) -> E2,
141 {
142 MapErr { backend: self, f }
143 }
144
145 fn with_codec<NewCodec>(self, codec: NewCodec) -> WithCodec<Self, NewCodec>
148 where
149 Self: Sized + BackendConfig,
150 NewCodec: Codec<Self::Args>,
151 {
152 WithCodec::new(self, codec)
153 }
154
155 fn poll_with_stream<S>(self, stream: S) -> PollWith<Self, StreamStrategy<S>>
157 where
158 Self: Sized,
159 S: Stream + Unpin + Send + 'static,
160 {
161 let strategy = StreamStrategy::new(stream);
162 PollWith::new(self, strategy)
163 }
164
165 #[cfg(feature = "sleep")]
167 fn poll_with_interval(self, duration: Duration) -> PollWith<Self, IntervalStrategy>
168 where
169 Self: Sized,
170 {
171 let strategy = IntervalStrategy::new(duration);
172 PollWith::new(self, strategy)
173 }
174
175 #[cfg(feature = "sleep")]
177 fn poll_with_backoff(
178 self,
179 interval: Duration,
180 config: BackoffConfig,
181 ) -> PollWith<Self, BackoffStrategy>
182 where
183 Self: Sized,
184 {
185 let strategy = IntervalStrategy::new(interval).with_backoff(config);
186 PollWith::new(self, strategy)
187 }
188
189 fn poll_with_strategy<S>(self, strategy: S) -> PollWith<Self, S>
191 where
192 Self: Sized,
193 S: PollStrategy,
194 {
195 PollWith::new(self, strategy)
196 }
197
198 #[cfg(feature = "tracing")]
199 fn instrumented(self, span: tracing::Span) -> Instrumented<Self>
201 where
202 Self: Sized,
203 {
204 Instrumented::new(self, span)
205 }
206
207 fn before_start<F, Fut>(self, f: F) -> BeforeStart<Self, Self::Error>
209 where
210 Self: Sized,
211 F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
212 Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
213 {
214 BeforeStart::new(self, f)
215 }
216
217 fn before_stop<F, Fut>(self, f: F) -> BeforeStop<Self, Self::Error>
219 where
220 Self: Sized,
221 F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
222 Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
223 {
224 BeforeStop::new(self, f)
225 }
226
227 fn after_start<F, Fut>(self, f: F) -> AfterStart<Self, Self::Error>
229 where
230 Self: Sized,
231 for<'c> F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
232 Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
233 {
234 AfterStart::new(self, f)
235 }
236
237 fn after_stop<F, Fut>(self, f: F) -> AfterStop<Self, Self::Error>
239 where
240 Self: Sized,
241 F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
242 Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
243 {
244 AfterStop::new(self, f)
245 }
246
247 fn interleave<S>(self, stream: S) -> Interleave<Self, S>
249 where
250 Self: Sized,
251 S: Stream<Item = Result<Self::Task, Self::Error>> + Unpin,
252 {
253 Interleave::new(self, stream)
254 }
255
256 fn wake_on_push(self) -> WakeOnPush<Self>
258 where
259 Self: Sized,
260 {
261 WakeOnPush::new(self)
262 }
263
264 #[cfg(feature = "shared")]
266 fn shared(self) -> Shared<Self>
267 where
268 Self: WireFormatBackend + Send,
269 Self::Codec: Clone,
270 {
271 Shared::new(self)
272 }
273}
274
275impl<B: Backend> BackendExt for B {}