apalis_file_storage/
sink.rs

1use std::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use apalis_core::task::{
7    Task,
8    task_id::{RandomId, TaskId},
9};
10use futures_sink::Sink;
11use serde_json::Value;
12
13use crate::{
14    Adapter, FileStorage, PendingChange, SyncPolicy, error::FileStorageError, util::RawTask,
15};
16
17impl<A: Adapter + Unpin, Args: Unpin> Sink<Task<Value>> for FileStorage<Args, A> {
18    type Error = FileStorageError<A>;
19
20    fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
21        Poll::Ready(Ok(()))
22    }
23
24    fn start_send(self: Pin<&mut Self>, mut item: Task<Value>) -> Result<(), Self::Error> {
25        let idempotency_key = item.idempotency_key().map(|a| a.to_owned());
26        if item.task_id().is_none() {
27            item = item
28                .into_builder()
29                .task_id(TaskId::from_string(RandomId::default()))
30                .build();
31        }
32
33        let entry = RawTask {
34            task_id: item.task_id().unwrap().clone(),
35            ctx: item.metadata().clone(),
36            args: item.args,
37            result: None,
38            idempotency_key,
39        };
40        let this = self.get_mut();
41        let tasks = this
42            .entries
43            .try_read()
44            .map_err(|_| FileStorageError::WouldBlockLock)?;
45        // Enforce idempotency
46        if entry.idempotency_key.is_some()
47            && tasks
48                .iter()
49                .any(|t| t.idempotency_key == entry.idempotency_key)
50        {
51            return Ok(());
52        }
53        drop(tasks);
54        let line = A::from_entry(&entry).map_err(FileStorageError::AdapterError)?;
55        let bytes = this
56            .adapter
57            .serialize(&line)
58            .map_err(FileStorageError::AdapterError)?;
59        let mut entries = this
60            .entries
61            .try_write()
62            .map_err(|_| FileStorageError::WouldBlockLock)?;
63        entries.push(entry); // optimistic in-memory update
64        drop(entries);
65        this.pending
66            .try_write()
67            .map_err(|_| FileStorageError::WouldBlockLock)?
68            .push(PendingChange::Append(bytes));
69
70        Ok(())
71    }
72
73    fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
74        match self.sync_policy.clone() {
75            SyncPolicy::Instant | SyncPolicy::Periodic(_) => self.flush()?,
76            SyncPolicy::Manual => {}
77        }
78        Poll::Ready(Ok(()))
79    }
80
81    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
82        self.poll_flush(cx)
83    }
84}
apalis_file_storage/sink.rs

apalis_file_storage/
sink.rs
