1use crate::backend::{BackendExt, TaskSink};
50use crate::error::BoxDynError;
51use crate::features_table;
52use crate::task::Task;
53use crate::{backend::Backend, backend::codec::Codec, worker::context::WorkerContext};
54use futures_sink::Sink;
55use futures_util::stream::{once, select};
56use futures_util::{SinkExt, Stream, TryStreamExt};
57use futures_util::{StreamExt, stream::BoxStream};
58use std::fmt;
59use std::fmt::Debug;
60use std::marker::PhantomData;
61use std::ops::{Deref, DerefMut};
62
63#[doc = features_table! {
65 setup = "{ unreachable!() }",
66 TaskSink => supported("Ability to push new tasks", false),
67 InheritsFeatures => limited("Inherits features from the underlying backend", false),
68}]
69pub struct Pipe<S, Into, Args, Ctx> {
70 pub(crate) from: S,
71 pub(crate) into: Into,
72 pub(crate) _req: PhantomData<(Args, Ctx)>,
73}
74
75impl<S: Clone, Into: Clone, Args, Ctx> Clone for Pipe<S, Into, Args, Ctx> {
76 fn clone(&self) -> Self {
77 Self {
78 from: self.from.clone(),
79 into: self.into.clone(),
80 _req: PhantomData,
81 }
82 }
83}
84
85impl<S, Into, Args, Ctx> Deref for Pipe<S, Into, Args, Ctx> {
86 type Target = Into;
87
88 fn deref(&self) -> &Self::Target {
89 &self.into
90 }
91}
92
93impl<S, Into, Args, Ctx> DerefMut for Pipe<S, Into, Args, Ctx> {
94 fn deref_mut(&mut self) -> &mut Self::Target {
95 &mut self.into
96 }
97}
98
99impl<S, Into, Args, Ctx> Pipe<S, Into, Args, Ctx> {
100 pub fn new(stream: S, backend: Into) -> Self {
102 Self {
103 from: stream,
104 into: backend,
105 _req: PhantomData,
106 }
107 }
108}
109
110impl<S: fmt::Debug, Into: fmt::Debug, Args, Ctx> fmt::Debug for Pipe<S, Into, Args, Ctx> {
111 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112 f.debug_struct("Pipe")
113 .field("inner", &self.from)
114 .field("into", &self.into)
115 .finish()
116 }
117}
118
119impl<Args, Ctx, S, TSink, Err> Backend for Pipe<S, TSink, Args, Ctx>
120where
121 S: Stream<Item = Result<Args, Err>> + Send + 'static,
122 TSink: Backend<Args = Args, Context = Ctx>
123 + BackendExt
124 + TaskSink<Args>
125 + Clone
126 + Unpin
127 + Send
128 + 'static
129 + Sink<Task<TSink::Compact, Ctx, TSink::IdType>>,
130 <TSink as Backend>::Error: std::error::Error + Send + Sync + 'static,
131 TSink::Beat: Send + 'static,
132 TSink::IdType: Send + Clone + 'static,
133 TSink::Stream: Send + 'static,
134 Args: Send + 'static,
135 Ctx: Send + 'static + Default,
136 Err: std::error::Error + Send + Sync + 'static,
137 <TSink as Sink<Task<TSink::Compact, Ctx, TSink::IdType>>>::Error:
138 std::error::Error + Send + Sync + 'static,
139 <<TSink as BackendExt>::Codec as Codec<Args>>::Error: std::error::Error + Send + Sync + 'static,
140 TSink::Compact: Send,
141{
142 type Args = Args;
143
144 type IdType = TSink::IdType;
145
146 type Context = Ctx;
147
148 type Stream = BoxStream<'static, Result<Option<Task<Args, Ctx, Self::IdType>>, PipeError>>;
149
150 type Layer = TSink::Layer;
151
152 type Beat = BoxStream<'static, Result<(), PipeError>>;
153
154 type Error = PipeError;
155
156 fn heartbeat(&self, worker: &WorkerContext) -> Self::Beat {
157 self.into
158 .heartbeat(worker)
159 .map_err(|e| PipeError::Inner(e.into()))
160 .boxed()
161 }
162
163 fn middleware(&self) -> Self::Layer {
164 self.into.middleware()
165 }
166
167 fn poll(self, worker: &WorkerContext) -> Self::Stream {
168 let mut sink = self
169 .into
170 .clone()
171 .sink_map_err(|e| PipeError::Inner(e.into()));
172
173 let mut sink_stream = self
174 .from
175 .map_err(|e| PipeError::Inner(e.into()))
176 .map_ok(|s| {
177 Task::new(s)
178 .try_map(|s| TSink::Codec::encode(&s).map_err(|e| PipeError::Inner(e.into())))
179 })
180 .map(|t| t.and_then(|t| t))
181 .boxed();
182
183 let sender_stream = self.into.poll(worker);
184 select(
185 once(async move {
186 let fut = sink.send_all(&mut sink_stream);
187 fut.await.map_err(|e| PipeError::Inner(e.into()))?;
188 Ok(None)
189 }),
190 sender_stream.map_err(|e| PipeError::Inner(e.into())),
191 )
192 .boxed()
193 }
194}
195
196pub trait PipeExt<B, Args, Ctx>
198where
199 Self: Sized,
200{
201 fn pipe_to(self, backend: B) -> Pipe<Self, B, Args, Ctx>;
203}
204
205impl<B, Args, Ctx, Err, S> PipeExt<B, Args, Ctx> for S
206where
207 S: Stream<Item = Result<Args, Err>> + Send + 'static,
208 <B as Backend>::Error: Into<BoxDynError> + Send + Sync + 'static,
209 B: Backend<Args = Args> + TaskSink<Args>,
210{
211 fn pipe_to(self, backend: B) -> Pipe<Self, B, Args, Ctx> {
212 Pipe::new(self, backend)
213 }
214}
215
216#[derive(Debug, thiserror::Error)]
218pub enum PipeError {
219 #[error("The inner stream provided a None")]
221 EmptyStream,
222 #[error("The inner stream error: {0}")]
224 Inner(BoxDynError),
225}
226
227#[cfg(test)]
228#[cfg(feature = "json")]
229mod tests {
230 use std::{io, time::Duration};
231
232 use futures_util::stream;
233
234 use crate::{
235 backend::json::JsonStorage,
236 error::BoxDynError,
237 worker::{
238 builder::WorkerBuilder, context::WorkerContext, ext::event_listener::EventListenerExt,
239 },
240 };
241
242 use super::*;
243
244 const ITEMS: u32 = 10;
245
246 #[tokio::test]
247 async fn basic_worker() {
248 let stm = stream::iter(0..ITEMS).map(Ok::<_, io::Error>);
249 let in_memory = JsonStorage::new_temp().unwrap();
250
251 let backend = Pipe::new(stm, in_memory);
252
253 async fn task(task: u32, ctx: WorkerContext) -> Result<(), BoxDynError> {
254 tokio::time::sleep(Duration::from_secs(1)).await;
255 if task == ITEMS - 1 {
256 ctx.stop().unwrap();
257 return Err("Graceful Exit".into());
258 }
259 Ok(())
260 }
261
262 let worker = WorkerBuilder::new("rango-tango")
263 .backend(backend)
264 .on_event(|_ctx, ev| {
265 println!("On Event = {ev:?}");
266 })
267 .build(task);
268 worker.run().await.unwrap();
269 }
270}