dge_gen/graph.rs
1use petgraph;
2use std::path::Path;
3
4pub use petgraph::graph::EdgeIndex;
5pub use petgraph::graph::NodeIndex;
6
7use super::generate;
8
9pub(crate) type PetGraph = petgraph::Graph<Node, Edge>;
10
11use crate::Result;
12
13/// A node represents the computation.
14///
15/// Every node has an unique name associated with it,
16/// which will be used as the file name of the generated file.
17#[derive(Clone, Debug)]
18pub(crate) enum Node {
19 /// A no-op node that indicates the start of the computation.
20 Start {
21 name: String,
22 },
23 /// A terminating node the indicates the termination of the computation.
24 Terminate {
25 name: String,
26 },
27 /// An aggregation node that consume one or more messages,
28 /// and outputs zero or more messages.
29 Aggregate {
30 name: String,
31 behaviour_module: String,
32 },
33 /// Duplicate the output of one node to multiple nodes.
34 FanOut {
35 name: String,
36 },
37 /// A user-provided handler that transform the input message into the output message.
38 UserHandler {
39 name: String,
40 behaviour_module: String,
41 },
42 /// A node that polls the incoming messages.
43 Poll {
44 name: String,
45 behaviour_module: String,
46 },
47}
48
49impl Node {
50 pub(crate) fn name(&self) -> String {
51 match self {
52 Node::Start { name, .. } => name.clone(),
53 Node::Terminate { name, .. } => name.clone(),
54 Node::Aggregate { name, .. } => name.clone(),
55 Node::FanOut { name, .. } => name.clone(),
56 Node::UserHandler { name, .. } => name.clone(),
57 Node::Poll { name, .. } => name.clone(),
58 }
59 }
60}
61
62/// An edge represents a RabbitMQ queue carrying a specific type of message.
63#[derive(Clone, Debug, Ord, PartialOrd, PartialEq, Eq, Hash)]
64pub(crate) struct Edge {
65 pub(crate) queue: String,
66 pub(crate) msg_type: String,
67 pub(crate) retry_interval_in_seconds: u32,
68}
69
70/// A computational graph where:
71///
72/// - edges represent messages of static types delivered between nodes via RabbitMQ queue
73/// - nodes represent computations that transform the input message into output message
74pub struct Graph {
75 pub(crate) g: PetGraph,
76 pub(crate) accept_failure: String,
77 pub(crate) type_error: String,
78}
79
80impl Graph {
81 /// Create a new computational graph.
82 ///
83 /// - `accept_failure : fn(Context, Error) -> Result<(), Error>`
84 /// - `type_error: Error`
85 ///
86 /// where:
87 ///
88 /// - `Error` is your global error type
89 /// (dge assumes that you use an global error type for your application)
90 /// - `Context` is a data type representing the current context,
91 /// this data type is defined by you,
92 /// and should be `From<T>` for all of your messages `T` carried by an edge
93 /// Different edges can carry different type of messages,
94 /// but they all have to be serializable and deserializable by serde_json,
95 /// this serialization requirement is strictly for convenience,
96 /// and can be removed if there is enough motivation.
97 pub fn new<S: Into<String>>(accept_failure: S, type_error: S) -> Graph {
98 Graph {
99 g: petgraph::Graph::new(),
100 accept_failure: accept_failure.into(),
101 type_error: type_error.into(),
102 }
103 }
104
105 /// Represent the start of the computation.
106 ///
107 /// Usually this is the first function been called to acquire a starting point
108 /// for later operations after the graph is created.
109 ///
110 /// Return a handle to the start node.
111 pub fn start<S: Into<String>>(&mut self, name: S) -> NodeIndex {
112 self.g.add_node(Node::Start { name: name.into() })
113 }
114
115 /// Read a message of type `type_input` from node `input` via RabbitMQ queue `queue`,
116 /// and process it with this node, which is named `name`,
117 /// using behaviour defined by `behaviour_module`,
118 /// retry after `retry_interval_in_seconds` if some error happened during the processing,
119 /// (transient or non-transient), and dge decides that the processing should be retried.
120 ///
121 /// Return a handle to the `handler` node.
122 ///
123 /// Internally this will create a new node for `handler`,
124 /// and a new edge from `input` to `handler` representing the underlying RabbitMQ queue `queue`.
125 pub fn process<S: Into<String>>(
126 &mut self,
127 input: NodeIndex,
128 queue: S,
129 type_input: S,
130 name: S,
131 behaviour_module: S,
132 retry_interval_in_seconds: u32,
133 ) -> NodeIndex {
134 let handler_node = Node::UserHandler {
135 name: name.into(),
136 behaviour_module: behaviour_module.into(),
137 };
138 let handler_node_i = self.g.add_node(handler_node);
139 let edge = Edge {
140 queue: queue.into(),
141 msg_type: type_input.into(),
142 retry_interval_in_seconds,
143 };
144 self.g.add_edge(input, handler_node_i, edge);
145 handler_node_i
146 }
147
148 /// Add a node that aggregate messages from `inputs` that belong to a single run,
149 /// and aggregate them for later consumption.
150 ///
151 /// `behaviour_module` defines how the input messages should be aggregated.
152 ///
153 /// Return a handle to the newly added node.
154 pub fn aggregate<S: Into<String>>(
155 &mut self,
156 inputs: Vec<NodeIndex>,
157 queue: S,
158 type_input: S,
159 name: S,
160 behaviour_module: S,
161 retry_interval_in_seconds: u32,
162 ) -> NodeIndex {
163 let type_input = type_input.into();
164 let wait_node_i = self.g.add_node(Node::Aggregate {
165 name: name.into(),
166 behaviour_module: behaviour_module.into(),
167 });
168 let queue = queue.into();
169 for input_i in inputs {
170 self.g.add_edge(
171 input_i,
172 wait_node_i,
173 Edge {
174 queue: queue.clone(),
175 msg_type: type_input.clone(),
176 retry_interval_in_seconds,
177 },
178 );
179 }
180 wait_node_i
181 }
182
183 /// Create a node that will copy messages of `input`
184 /// to all outgoing edges of the newly created node.
185 ///
186 /// Return a handle to the newly created node
187 pub fn fan_out<S: Into<String>>(
188 &mut self,
189 input: NodeIndex,
190 queue: S,
191 type_input: S,
192 name: S,
193 retry_interval_in_seconds: u32,
194 ) -> NodeIndex {
195 let fan_out_i = self.g.add_node(Node::FanOut { name: name.into() });
196 self.g.add_edge(
197 input,
198 fan_out_i,
199 Edge {
200 queue: queue.into(),
201 msg_type: type_input.into(),
202 retry_interval_in_seconds,
203 },
204 );
205
206 fan_out_i
207 }
208
209 /// Add a node that polls some external system using the input message as arguments.
210 ///
211 /// `behaviour_module` defines the function that will perform the polling,
212 /// this nodes provides scheduling for the actual polling function.
213 ///
214 /// For example this can be used to query an third-party service for the availability
215 /// of resource corresponding the input messages.
216 pub fn poll<S: Into<String>>(
217 &mut self,
218 input: NodeIndex,
219 queue: S,
220 type_input: S,
221 name: S,
222 behaviour_module: S,
223 retry_interval_in_seconds: u32,
224 ) -> NodeIndex {
225 let poll_node = Node::Poll {
226 name: name.into(),
227 behaviour_module: behaviour_module.into(),
228 };
229 let poll_node_i = self.g.add_node(poll_node);
230 let edge = Edge {
231 queue: queue.into(),
232 msg_type: type_input.into(),
233 retry_interval_in_seconds,
234 };
235 self.g.add_edge(input, poll_node_i, edge);
236 poll_node_i
237 }
238
239 /// An no-op node that terminates the computation.
240 pub fn terminate<S: Into<String>>(&mut self, input: NodeIndex, queue: S, type_input: S, name: S, retry_interval_in_seconds: u32) -> () {
241 let terminate_node = self.g.add_node(Node::Terminate { name: name.into() });
242 self.g.add_edge(input, terminate_node, Edge {
243 queue: queue.into(),
244 msg_type: type_input.into(),
245 retry_interval_in_seconds
246 });
247 }
248
249
250 /// Generate code represented by the graph.
251 ///
252 /// - the generated code will be written to `output_dir`
253 /// - `get_rmq_uri` is used to get an url to connect to the RabbitMQ server
254 /// - `work_exchange` and `retry_exchange` are direct RabbitMQ exchanges,
255 /// for every edge in the graph, there will be a queue bound to `work_exchange`
256 /// to deliver the message, and a retry queue bound to `retry_exchange` to handle the retry,
257 /// (the retry is backed by RabbitMQ's dead lettering mechanism)
258 /// - if the `init_input_queue` is true, then queues originated from starts node
259 /// are also declared by the `init-exchanges-and-queues` subcommand.
260 /// - `init_output_queue` controls the initialization of queues leading to termination nodes
261 /// - `main_init` is a function that will be run prior to the start of the computation,
262 /// this can be used for things like setting up the logger
263 pub fn generate<P: AsRef<Path>, S: AsRef<str>>(
264 self,
265 output_dir: P,
266 get_rmq_uri: S,
267 work_exchange: S,
268 retry_exchange: S,
269 retry_queue_prefix: S,
270 retry_queue_suffix: S,
271 init_input_queue: bool,
272 init_output_queue: bool,
273 main_init: S,
274 ) -> Result<()> {
275 let rmq_options = generate::graph::RmqOptions {
276 get_rmq_uri: get_rmq_uri.as_ref().into(),
277 work_exchange: work_exchange.as_ref().into(),
278 retry_exchange: retry_exchange.as_ref().into(),
279 retry_queue_prefix: retry_queue_prefix.as_ref().into(),
280 retry_queue_suffix: retry_queue_suffix.as_ref().into(),
281 };
282 generate::graph::generate(self, output_dir, rmq_options, init_input_queue, init_output_queue, main_init)
283 }
284}