apalis_core/backend/ext/
shared.rs

1use std::collections::VecDeque;
2use std::pin::Pin;
3use std::sync::Arc;
4use std::sync::atomic::AtomicUsize;
5use std::task::{Context, Poll, Waker};
6
7use crate::backend::*;
8use crate::task::Task;
9use crate::worker::context::WorkerContext;
10use dashmap::DashMap;
11use futures_core::ready;
12use futures_sink::Sink;
13use futures_util::lock::Mutex;
14use futures_util::{FutureExt, SinkExt, StreamExt};
15
16#[derive(Debug)]
17struct Inner<B> {
18    backend: Mutex<B>,
19    wakers: DashMap<usize, Waker>,
20    next_key: std::sync::atomic::AtomicUsize,
21}
22
23/// Create a cloneable handle to the inner backend where all handles are clone.
24#[derive(Debug)]
25pub struct Shared<B: WireFormatBackend> {
26    inner: Arc<Inner<B>>,
27    waker_key: usize,
28    codec: B::Codec,
29    sink: VecDeque<Task<B::Compact>>,
30}
31
32impl<B> Clone for Shared<B>
33where
34    B::Codec: Clone,
35    B: WireFormatBackend,
36{
37    fn clone(&self) -> Self {
38        let inner = self.inner.clone();
39        let waker_key = inner
40            .next_key
41            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
42        Self {
43            inner,
44            waker_key,
45            codec: self.codec.clone(),
46            sink: Default::default(),
47        }
48    }
49}
50
51impl<B: WireFormatBackend> Shared<B>
52where
53    B::Codec: Clone,
54{
55    /// Build a new Sharable backend
56    pub fn new(backend: B) -> Self {
57        let codec = backend.codec().clone();
58        Self {
59            inner: Inner {
60                backend: Mutex::new(backend),
61                next_key: AtomicUsize::new(0),
62                wakers: Default::default(),
63            }
64            .into(),
65            waker_key: 0,
66            codec,
67            sink: VecDeque::new(),
68        }
69    }
70}
71
72impl<B: WireFormatBackend> Shared<B> {
73    fn register_waker(inner: &Inner<B>, key: usize, cx: &Context<'_>) {
74        let res = &inner.wakers;
75        res.insert(key, cx.waker().clone());
76    }
77
78    fn wake_others(inner: &Inner<B>, except: usize) {
79        let wakers = &inner.wakers;
80        wakers.retain(|key, waker| {
81            if *key != except {
82                waker.wake_by_ref();
83            }
84            true
85        });
86    }
87}
88
89impl<B> Backend for Shared<B>
90where
91    B: Backend + WireFormatBackend,
92{
93    type Task = B::Task;
94    type Error = B::Error;
95
96    fn poll_ready(
97        &mut self,
98        cx: &mut Context<'_>,
99        worker: &WorkerContext,
100    ) -> Poll<Result<(), Self::Error>> {
101        if let Some(mut guard) = self.inner.backend.try_lock() {
102            let result = guard.poll_ready(cx, worker);
103            drop(guard);
104
105            if result.is_ready() {
106                Self::wake_others(&self.inner, self.waker_key);
107            }
108            result
109        } else {
110            Self::register_waker(&self.inner, self.waker_key, cx);
111            Poll::Pending
112        }
113    }
114
115    fn poll_next(
116        &mut self,
117        cx: &mut Context<'_>,
118        worker: &WorkerContext,
119    ) -> Poll<Option<Result<Self::Task, Self::Error>>> {
120        let inner = &self.inner;
121
122        if let Some(mut guard) = inner.backend.try_lock() {
123            let result = guard.poll_next(cx, worker);
124            drop(guard);
125            Self::wake_others(inner, self.waker_key);
126            result
127        } else {
128            Self::register_waker(inner, self.waker_key, cx);
129            Poll::Pending
130        }
131    }
132
133    fn poll_close(
134        &mut self,
135        cx: &mut Context<'_>,
136        worker: &WorkerContext,
137    ) -> Poll<Result<(), Self::Error>> {
138        let inner = self.inner.as_ref();
139
140        if let Some(mut guard) = inner.backend.try_lock() {
141            let result = guard.poll_close(cx, worker);
142            drop(guard);
143            if result.is_ready() {
144                inner.wakers.remove(&self.waker_key);
145                Self::wake_others(inner, self.waker_key);
146            } else {
147                Self::wake_others(inner, self.waker_key);
148            }
149            result
150        } else {
151            Self::register_waker(inner, self.waker_key, cx);
152            Poll::Pending
153        }
154    }
155}
156
157impl<B: WireFormatBackend> WireFormatBackend for Shared<B> {
158    type Codec = B::Codec;
159
160    type Compact = B::Compact;
161
162    fn codec(&self) -> &Self::Codec {
163        &self.codec
164    }
165}
166
167impl<B, Err> Sink<Task<B::Compact>> for Shared<B>
168where
169    B: Backend<Error = Err> + Sink<Task<B::Compact>, Error = Err> + Unpin,
170    B: WireFormatBackend,
171    B::Codec: Unpin,
172    B::Compact: Unpin,
173{
174    type Error = Err;
175    fn start_send(self: Pin<&mut Self>, item: Task<B::Compact>) -> Result<(), Self::Error> {
176        let this = self.get_mut();
177
178        if let Some(mut guard) = this.inner.backend.try_lock() {
179            let result = guard.start_send_unpin(item);
180            drop(guard);
181            if result.is_ok() {
182                Self::wake_others(&this.inner, this.waker_key);
183            }
184            result
185        } else {
186            this.sink.push_back(item);
187            Ok(())
188        }
189    }
190    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
191        let this = self.get_mut();
192        let mut guard = ready!(this.inner.backend.lock().poll_unpin(cx));
193        while !this.sink.is_empty() {
194            ready!(guard.poll_ready_unpin(cx))?;
195            let item = this.sink.pop_front().unwrap();
196            guard.start_send_unpin(item)?;
197        }
198
199        guard.poll_ready_unpin(cx)
200    }
201    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
202        let mut guard = ready!(self.inner.backend.lock().poll_unpin(cx));
203        guard.poll_flush_unpin(cx)
204    }
205    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
206        let mut guard = ready!(self.inner.backend.lock().poll_unpin(cx));
207        guard.poll_close_unpin(cx)
208    }
209}
210
211// delegate_deref!(Shared<B>, inner);
212
213impl<B> BackendConfig for Shared<B>
214where
215    B: BackendConfig + WireFormatBackend,
216{
217    type Id = B::Id;
218    type Args = B::Args;
219    type Kind = B::Kind;
220    type Config = B::Config;
221    type Layer = B::Layer;
222    fn config(&self) -> &Self::Config {
223        unreachable!("Dont call config on shared")
224    }
225    fn middleware(&mut self, worker: &mut WorkerContext) -> Self::Layer {
226        self.inner.backend.try_lock().unwrap().middleware(worker)
227    }
228}
229
230impl<B> FetchById for Shared<B>
231where
232    B: FetchById,
233    B::Task: Send,
234    B: Backend + WireFormatBackend + Send,
235    B::Codec: Send,
236    B::Compact: Send,
237{
238    async fn fetch_by_id(
239        &mut self,
240        task_id: &crate::task::task_id::TaskId,
241    ) -> Result<Option<B::Task>, Self::Error> {
242        self.inner.backend.lock().await.fetch_by_id(task_id).await
243    }
244}
245impl<B> Update for Shared<B>
246where
247    B: Update + Send,
248    B::Task: Send,
249    B: Backend + WireFormatBackend + Send,
250    B::Codec: Send,
251    B::Compact: Send,
252{
253    async fn update(&mut self, task: Self::Task) -> Result<(), Self::Error> {
254        self.inner.backend.lock().await.update(task).await
255    }
256}
257impl<B> Reschedule for Shared<B>
258where
259    B: Reschedule,
260    B::Task: Send,
261    B: Backend + WireFormatBackend + Send,
262    B::Codec: Send,
263    B::Compact: Send,
264{
265    async fn reschedule(
266        &mut self,
267        task: Self::Task,
268        wait: std::time::Duration,
269    ) -> Result<(), Self::Error> {
270        self.inner.backend.lock().await.reschedule(task, wait).await
271    }
272}
273impl<B> Vacuum for Shared<B>
274where
275    B: Vacuum,
276    B: Backend + WireFormatBackend + Send,
277    B::Codec: Send,
278    B::Compact: Send,
279{
280    async fn vacuum(&mut self) -> Result<usize, Self::Error> {
281        self.inner.backend.lock().await.vacuum().await
282    }
283}
284impl<B> ResumeById for Shared<B>
285where
286    B: ResumeById,
287    B: Backend + WireFormatBackend + Send,
288    B::Codec: Send,
289    B::Compact: Send,
290{
291    async fn resume_by_id(&mut self, id: TaskId) -> Result<bool, Self::Error> {
292        self.inner.backend.lock().await.resume_by_id(id).await
293    }
294}
295impl<B> ResumeAbandoned for Shared<B>
296where
297    B: ResumeAbandoned,
298    B: Backend + WireFormatBackend + Send,
299    B::Codec: Send,
300    B::Compact: Send,
301{
302    async fn resume_abandoned(&mut self) -> Result<usize, Self::Error> {
303        self.inner.backend.lock().await.resume_abandoned().await
304    }
305}
306impl<B> RegisterWorker for Shared<B>
307where
308    B: RegisterWorker,
309    B: Backend + WireFormatBackend + Send,
310    B::Codec: Send,
311    B::Compact: Send,
312{
313    async fn register_worker(&mut self, worker_id: String) -> Result<(), Self::Error> {
314        self.inner
315            .backend
316            .lock()
317            .await
318            .register_worker(worker_id)
319            .await
320    }
321}
322impl<Output, B> WaitForCompletion<Output> for Shared<B>
323where
324    B: WaitForCompletion<Output> + Sync + 'static,
325    Output: Send + 'static,
326    B: Backend + WireFormatBackend + Send,
327    B::Codec: Sync,
328    B::Codec: Send,
329    B::Compact: Send + Sync,
330{
331    type ResultStream =
332        futures_core::stream::BoxStream<'static, Result<TaskResult<Output>, Self::Error>>;
333    fn wait_for(&self, task_ids: impl IntoIterator<Item = TaskId>) -> Self::ResultStream {
334        let inner = self.inner.clone();
335        let task_ids: Vec<_> = task_ids.into_iter().collect();
336        futures_util::stream::once(async move {
337            let backend = inner.backend.lock().await;
338            backend.wait_for(task_ids)
339        })
340        .flatten()
341        .boxed()
342    }
343    async fn check_status(
344        &self,
345        task_ids: impl IntoIterator<Item = TaskId> + Send,
346    ) -> Result<Vec<TaskResult<Output>>, Self::Error> {
347        self.inner.backend.lock().await.check_status(task_ids).await
348    }
349}
350
351impl<B: WireFormatBackend> Drop for Shared<B> {
352    fn drop(&mut self) {
353        if let Some((_key, waker)) = self.inner.wakers.remove(&self.waker_key) {
354            waker.wake();
355        }
356    }
357}
apalis_core/backend/ext/shared.rs

apalis_core/backend/ext/
shared.rs
