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 shared::Shared,
36 wake_on_push::WakeOnPush,
37 with_codec::WithCodec,
38 },
39 },
40 error::BoxDynError,
41 task::Task,
42 worker::context::WorkerContext,
43};
44
45#[cfg(feature = "sleep")]
46use crate::backend::ext::poll_strategy::{BackoffConfig, BackoffStrategy, IntervalStrategy};
47
48#[macro_use]
49pub mod delegate;
50pub mod inspect_err;
52
53#[cfg(feature = "tracing")]
55pub mod instrument;
56pub mod interleave;
58
59pub mod map_err;
61pub mod pipe;
62pub mod poll_strategy;
63pub mod wake_on_push;
65pub mod with_codec;
66
67pub mod lifecycle;
68
69pub mod shared;
71
72#[derive(Debug, thiserror::Error)]
74#[non_exhaustive]
75pub enum PollNextArgsError<B: Backend> {
76 #[error("backend error: {0}")]
78 BackendError(B::Error),
79 #[error("failed to decode task args: {0}")]
81 DecodeError(BoxDynError),
82}
83
84pub trait BackendExt: Backend {
86 #[allow(clippy::type_complexity)]
89 fn poll_next_args(
90 &mut self,
91 cx: &mut Context<'_>,
92 worker: &WorkerContext,
93 ) -> Poll<Option<Result<Task<Self::Args>, PollNextArgsError<Self>>>>
94 where
95 Self: Sized + BackendConfig + WireFormatBackend + Backend<Task = Task<Self::Compact>>,
96 Self::Codec: Codec<Self::Args, Compact = Self::Compact>,
97 <Self::Codec as Codec<Self::Args>>::Error: std::error::Error + Send + Sync + 'static,
98 {
99 let next = self.poll_next(cx, worker);
100 let codec = self.codec();
101 next.map(move |item| match item {
102 Some(Ok(task)) => {
103 let task = task.try_map_args(|compact| codec.decode(&compact));
104 Some(task.map_err(|e| PollNextArgsError::DecodeError(e.into())))
105 }
106 Some(Err(e)) => Some(Err(PollNextArgsError::BackendError(e))),
107 None => None,
108 })
109 }
110
111 fn pipe_to<Dst>(self, backend: Dst) -> Pipe<Dst, Self>
117 where
118 Self: Sized,
119 {
120 Pipe::new(self, backend)
121 }
122
123 fn inspect_err<F>(self, f: F) -> InspectErr<Self, F>
125 where
126 Self: Sized,
127 F: Fn(&Self::Error),
128 {
129 InspectErr { backend: self, f }
130 }
131
132 fn map_err<F, E2>(self, f: F) -> MapErr<Self, F>
135 where
136 Self: Sized,
137 F: Fn(Self::Error) -> E2,
138 {
139 MapErr { backend: self, f }
140 }
141
142 fn with_codec<NewCodec>(self, codec: NewCodec) -> WithCodec<Self, NewCodec>
145 where
146 Self: Sized + BackendConfig,
147 NewCodec: Codec<Self::Args>,
148 {
149 WithCodec::new(self, codec)
150 }
151
152 fn poll_with_stream<S>(self, stream: S) -> PollWith<Self, StreamStrategy<S>>
154 where
155 Self: Sized,
156 S: Stream + Unpin + Send + 'static,
157 {
158 let strategy = StreamStrategy::new(stream);
159 PollWith::new(self, strategy)
160 }
161
162 #[cfg(feature = "sleep")]
164 fn poll_with_interval(self, duration: Duration) -> PollWith<Self, IntervalStrategy>
165 where
166 Self: Sized,
167 {
168 let strategy = IntervalStrategy::new(duration);
169 PollWith::new(self, strategy)
170 }
171
172 #[cfg(feature = "sleep")]
174 fn poll_with_backoff(
175 self,
176 interval: Duration,
177 config: BackoffConfig,
178 ) -> PollWith<Self, BackoffStrategy>
179 where
180 Self: Sized,
181 {
182 let strategy = IntervalStrategy::new(interval).with_backoff(config);
183 PollWith::new(self, strategy)
184 }
185
186 fn poll_with_strategy<S>(self, strategy: S) -> PollWith<Self, S>
188 where
189 Self: Sized,
190 S: PollStrategy,
191 {
192 PollWith::new(self, strategy)
193 }
194
195 #[cfg(feature = "tracing")]
196 fn instrumented(self, span: tracing::Span) -> Instrumented<Self>
198 where
199 Self: Sized,
200 {
201 Instrumented::new(self, span)
202 }
203
204 fn before_start<F, Fut>(self, f: F) -> BeforeStart<Self, Self::Error>
206 where
207 Self: Sized,
208 F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
209 Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
210 {
211 BeforeStart::new(self, f)
212 }
213
214 fn before_stop<F, Fut>(self, f: F) -> BeforeStop<Self, Self::Error>
216 where
217 Self: Sized,
218 F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
219 Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
220 {
221 BeforeStop::new(self, f)
222 }
223
224 fn after_start<F, Fut>(self, f: F) -> AfterStart<Self, Self::Error>
226 where
227 Self: Sized,
228 for<'c> F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
229 Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
230 {
231 AfterStart::new(self, f)
232 }
233
234 fn after_stop<F, Fut>(self, f: F) -> AfterStop<Self, Self::Error>
236 where
237 Self: Sized,
238 F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
239 Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
240 {
241 AfterStop::new(self, f)
242 }
243
244 fn interleave<S>(self, stream: S) -> Interleave<Self, S>
246 where
247 Self: Sized,
248 S: Stream<Item = Result<Self::Task, Self::Error>> + Unpin,
249 {
250 Interleave::new(self, stream)
251 }
252
253 fn wake_on_push(self) -> WakeOnPush<Self>
255 where
256 Self: Sized,
257 {
258 WakeOnPush::new(self)
259 }
260
261 fn shared(self) -> Shared<Self>
263 where
264 Self: WireFormatBackend + Send,
265 Self::Codec: Clone,
266 {
267 Shared::new(self)
268 }
269}
270
271impl<B: Backend> BackendExt for B {}