apalis-workflow 1.0.0

A flexible and composable task workflow engine for rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

use std::{fmt::Display, str::FromStr};

use apalis_core::{
    backend::{Backend, BackendConfig, TaskSinkError, WireFormatBackend, codec::Codec},
    error::BoxDynError,
    task::{Task, builder::TaskBuilder, task_id::GenerateId},
};
use futures_sink::Sink;
use petgraph::graph::NodeIndex;

use crate::{
    graph::{GraphFlowContext, decode::GraphCodec},
    sequential::WorkflowContext,
};

/// Extension trait for pushing tasks into a workflow
pub trait WorkflowSink<Args>: WireFormatBackend + Backend + Sized
where
    Self::Codec: Codec<Args, Compact = Self::Compact>,
{
    /// Push a single task into the workflow sink at the start
    fn push_start(
        &mut self,
        args: Args,
    ) -> impl Future<Output = Result<(), TaskSinkError<Self::Error>>> + Send;

    /// Push a single task into the workflow sink at the start
    fn start_fan_out(
        &mut self,
        args: Args,
    ) -> impl Future<Output = Result<(), TaskSinkError<Self::Error>>> + Send
    where
        Args: GraphCodec<Self>,
        Args::Error: std::error::Error + Send + Sync + 'static;

    /// Push a step into the workflow sink at the specified index
    ///
    /// This is a helper method for pushing tasks into the workflow sink
    /// with the appropriate workflow context metadata.
    /// Ideally, this should be used internally by the workflow executor
    /// rather than being called directly.
    fn push_step(
        &mut self,
        args: Args,
        index: usize,
    ) -> impl Future<Output = Result<(), TaskSinkError<Self::Error>>> + Send;

    /// Push a node into the workflow sink at the specified index
    ///
    /// This is a helper method for pushing tasks into the workflow sink
    /// with the appropriate Graph flow context metadata.
    /// Ideally, this should be used internally by the Graph executor
    /// rather than being called directly.
    fn push_node(
        &mut self,
        node: Args,
        index: NodeIndex,
    ) -> impl Future<Output = Result<(), TaskSinkError<Self::Error>>> + Send;
}

impl<S: Send, Args: Send, Compact, Err> WorkflowSink<Args> for S
where
    S: Sink<Task<Compact>, Error = Err>
        + Backend<Error = Err>
        + WireFormatBackend<Compact = Compact>
        + BackendConfig
        + Unpin,
    S::Id: GenerateId + Send + Sync + FromStr + Display,
    S::Codec: Codec<Args, Compact = Compact>,
    Err: std::error::Error + Send + Sync + 'static,
    <S::Codec as Codec<Args>>::Error: Into<BoxDynError> + Send + Sync + 'static,
    Compact: Send + 'static,
    <S::Id as FromStr>::Err: std::error::Error + Send + Sync + 'static,
{
    async fn push_start(&mut self, args: Args) -> Result<(), TaskSinkError<Self::Error>> {
        use futures_util::SinkExt;

        let codec = self.codec();
        let task_id = S::Id::generate();
        let compact =
            S::Codec::encode(codec, &args).map_err(|e| TaskSinkError::CodecError(e.into()))?;
        let task = TaskBuilder::new(compact).task_id(task_id.clone()).build();
        self.send(task)
            .await
            .map_err(|e| TaskSinkError::PushError(e))
    }

    async fn start_fan_out(&mut self, args: Args) -> Result<(), TaskSinkError<Self::Error>>
    where
        Args: GraphCodec<Self>,
        Args::Error: std::error::Error + Send + Sync + 'static,
    {
        use futures_util::SinkExt;
        let task_id = S::Id::generate();
        let codec = self.codec();
        let compact = Args::encode(args, codec).map_err(|e| TaskSinkError::CodecError(e.into()))?;
        let task = TaskBuilder::new(compact).task_id(task_id.clone()).build();
        self.send(task)
            .await
            .map_err(|e| TaskSinkError::PushError(e))
    }

    async fn push_step(
        &mut self,
        step: Args,
        index: usize,
    ) -> Result<(), TaskSinkError<Self::Error>> {
        use futures_util::SinkExt;
        let task_id = S::Id::generate();
        let codec = self.codec();
        let compact =
            S::Codec::encode(codec, &step).map_err(|e| TaskSinkError::CodecError(e.into()))?;
        let task = TaskBuilder::new(compact)
            .metadata(&WorkflowContext { step_index: index })
            .task_id(task_id.clone())
            .build();
        self.send(task)
            .await
            .map_err(|e| TaskSinkError::PushError(e))
    }

    async fn push_node(
        &mut self,
        node: Args,
        index: NodeIndex,
    ) -> Result<(), TaskSinkError<Self::Error>> {
        use futures_util::SinkExt;
        let task_id = S::Id::generate();
        let codec = self.codec();
        let compact =
            S::Codec::encode(codec, &node).map_err(|e| TaskSinkError::CodecError(e.into()))?;
        let task = TaskBuilder::new(compact)
            .metadata(&GraphFlowContext {
                current_node: index,
                completed_nodes: Default::default(),
                current_position: index.index(),
                is_initial: true,
                node_task_ids: Default::default(),
                prev_node: None,
                root_task_id: Some(task_id.clone()),
            })
            .task_id(task_id.clone())
            .build();
        self.send(task)
            .await
            .map_err(|e| TaskSinkError::PushError(e))
    }
}