apalis_workflow/
sink.rs

1use std::{fmt::Display, str::FromStr};
2
3use apalis_core::{
4    backend::{Backend, BackendConfig, TaskSinkError, WireFormatBackend, codec::Codec},
5    error::BoxDynError,
6    task::{Task, builder::TaskBuilder, task_id::GenerateId},
7};
8use futures_sink::Sink;
9use petgraph::graph::NodeIndex;
10
11use crate::{
12    graph::{GraphFlowContext, decode::GraphCodec},
13    sequential::WorkflowContext,
14};
15
16/// Extension trait for pushing tasks into a workflow
17pub trait WorkflowSink<Args>: WireFormatBackend + Backend + Sized
18where
19    Self::Codec: Codec<Args, Compact = Self::Compact>,
20{
21    /// Push a single task into the workflow sink at the start
22    fn push_start(
23        &mut self,
24        args: Args,
25    ) -> impl Future<Output = Result<(), TaskSinkError<Self::Error>>> + Send;
26
27    /// Push a single task into the workflow sink at the start
28    fn start_fan_out(
29        &mut self,
30        args: Args,
31    ) -> impl Future<Output = Result<(), TaskSinkError<Self::Error>>> + Send
32    where
33        Args: GraphCodec<Self>,
34        Args::Error: std::error::Error + Send + Sync + 'static;
35
36    /// Push a step into the workflow sink at the specified index
37    ///
38    /// This is a helper method for pushing tasks into the workflow sink
39    /// with the appropriate workflow context metadata.
40    /// Ideally, this should be used internally by the workflow executor
41    /// rather than being called directly.
42    fn push_step(
43        &mut self,
44        args: Args,
45        index: usize,
46    ) -> impl Future<Output = Result<(), TaskSinkError<Self::Error>>> + Send;
47
48    /// Push a node into the workflow sink at the specified index
49    ///
50    /// This is a helper method for pushing tasks into the workflow sink
51    /// with the appropriate Graph flow context metadata.
52    /// Ideally, this should be used internally by the Graph executor
53    /// rather than being called directly.
54    fn push_node(
55        &mut self,
56        node: Args,
57        index: NodeIndex,
58    ) -> impl Future<Output = Result<(), TaskSinkError<Self::Error>>> + Send;
59}
60
61impl<S: Send, Args: Send, Compact, Err> WorkflowSink<Args> for S
62where
63    S: Sink<Task<Compact>, Error = Err>
64        + Backend<Error = Err>
65        + WireFormatBackend<Compact = Compact>
66        + BackendConfig
67        + Unpin,
68    S::Id: GenerateId + Send + Sync + FromStr + Display,
69    S::Codec: Codec<Args, Compact = Compact>,
70    Err: std::error::Error + Send + Sync + 'static,
71    <S::Codec as Codec<Args>>::Error: Into<BoxDynError> + Send + Sync + 'static,
72    Compact: Send + 'static,
73    <S::Id as FromStr>::Err: std::error::Error + Send + Sync + 'static,
74{
75    async fn push_start(&mut self, args: Args) -> Result<(), TaskSinkError<Self::Error>> {
76        use futures_util::SinkExt;
77
78        let codec = self.codec();
79        let task_id = S::Id::generate();
80        let compact =
81            S::Codec::encode(codec, &args).map_err(|e| TaskSinkError::CodecError(e.into()))?;
82        let task = TaskBuilder::new(compact).task_id(task_id.clone()).build();
83        self.send(task)
84            .await
85            .map_err(|e| TaskSinkError::PushError(e))
86    }
87
88    async fn start_fan_out(&mut self, args: Args) -> Result<(), TaskSinkError<Self::Error>>
89    where
90        Args: GraphCodec<Self>,
91        Args::Error: std::error::Error + Send + Sync + 'static,
92    {
93        use futures_util::SinkExt;
94        let task_id = S::Id::generate();
95        let codec = self.codec();
96        let compact = Args::encode(args, codec).map_err(|e| TaskSinkError::CodecError(e.into()))?;
97        let task = TaskBuilder::new(compact).task_id(task_id.clone()).build();
98        self.send(task)
99            .await
100            .map_err(|e| TaskSinkError::PushError(e))
101    }
102
103    async fn push_step(
104        &mut self,
105        step: Args,
106        index: usize,
107    ) -> Result<(), TaskSinkError<Self::Error>> {
108        use futures_util::SinkExt;
109        let task_id = S::Id::generate();
110        let codec = self.codec();
111        let compact =
112            S::Codec::encode(codec, &step).map_err(|e| TaskSinkError::CodecError(e.into()))?;
113        let task = TaskBuilder::new(compact)
114            .metadata(&WorkflowContext { step_index: index })
115            .task_id(task_id.clone())
116            .build();
117        self.send(task)
118            .await
119            .map_err(|e| TaskSinkError::PushError(e))
120    }
121
122    async fn push_node(
123        &mut self,
124        node: Args,
125        index: NodeIndex,
126    ) -> Result<(), TaskSinkError<Self::Error>> {
127        use futures_util::SinkExt;
128        let task_id = S::Id::generate();
129        let codec = self.codec();
130        let compact =
131            S::Codec::encode(codec, &node).map_err(|e| TaskSinkError::CodecError(e.into()))?;
132        let task = TaskBuilder::new(compact)
133            .metadata(&GraphFlowContext {
134                current_node: index,
135                completed_nodes: Default::default(),
136                current_position: index.index(),
137                is_initial: true,
138                node_task_ids: Default::default(),
139                prev_node: None,
140                root_task_id: Some(task_id.clone()),
141            })
142            .task_id(task_id.clone())
143            .build();
144        self.send(task)
145            .await
146            .map_err(|e| TaskSinkError::PushError(e))
147    }
148}
apalis_workflow/sink.rs

apalis_workflow/
sink.rs
