bpm_engine_runtime/
token_arrived_handler.rs1use async_trait::async_trait;
5use bpm_engine_core::{payloads, EngineEvent, InstanceState, NodeType, TokenStatus};
6use std::collections::{HashMap, HashSet};
7use std::sync::Mutex;
8use tracing::{debug, warn};
9
10use super::handler::{EngineContext, EventHandler};
11use super::transition::{evaluate_exclusive_gateway, move_token, move_token_with_group};
12
13pub struct TokenArrivedHandler {
14 join_state: Mutex<HashMap<String, (usize, HashSet<String>)>>,
15}
16
17impl TokenArrivedHandler {
18 pub fn new() -> Self {
19 TokenArrivedHandler {
20 join_state: Mutex::new(HashMap::new()),
21 }
22 }
23}
24
25impl Default for TokenArrivedHandler {
26 fn default() -> Self {
27 Self::new()
28 }
29}
30
31#[async_trait]
32impl EventHandler for TokenArrivedHandler {
33 async fn handle(&self, event: &EngineEvent, ctx: &mut EngineContext) -> Vec<EngineEvent> {
34 let EngineEvent::TokenArrived(e) = event else {
35 return vec![];
36 };
37 let Some(process_store) = ctx.process_store.as_ref() else {
38 return vec![];
39 };
40 let Some(process_def_store) = ctx.process_def_store.as_ref() else {
41 return vec![];
42 };
43 let Ok(Some(mut instance)) = process_store.load(&e.instance_id).await else {
44 return vec![];
45 };
46 let Ok(Some(def)) = process_def_store.load(&instance.process_def_id).await else {
47 return vec![];
48 };
49 let Some(token_idx) = instance.tokens.iter().position(|t| t.id == e.token_id) else {
50 return vec![];
51 };
52 let token = &instance.tokens[token_idx];
53 if let Some(tr) = ctx.token_store.as_ref() {
54 if !tr
55 .claim_token(&e.instance_id, &token.id, token.version)
56 .await
57 .unwrap_or(false)
58 {
59 warn!(instance_id = %e.instance_id, token_id = %token.id, "token claim failed (CAS)");
60 return vec![];
61 }
62 instance.tokens[token_idx].status = TokenStatus::Executing;
63 instance.tokens[token_idx].version += 1;
64 }
65 debug!(instance_id = %e.instance_id, token_id = %e.token_id, node_id = %e.node_id, "token arrived");
66 let node = match def.nodes.get(e.node_id.as_str()) {
67 Some(n) => n,
68 None => return vec![],
69 };
70
71 let mut out = vec![];
72
73 match &node.node_type {
74 NodeType::Start => {
75 instance.tokens[token_idx].status = TokenStatus::Waiting;
76 let new_tokens = move_token(node);
77 for t in &new_tokens {
78 out.push(EngineEvent::TokenArrived(payloads::TokenArrived {
79 instance_id: e.instance_id.clone(),
80 token_id: t.id.clone(),
81 node_id: t.node_id.clone(),
82 }));
83 }
84 instance.tokens.extend(new_tokens);
85 }
86 NodeType::ServiceTask(service) => {
87 service(&mut instance);
88 instance.tokens[token_idx].status = TokenStatus::Waiting;
89 let new_tokens = move_token(node);
90 for t in &new_tokens {
91 out.push(EngineEvent::TokenArrived(payloads::TokenArrived {
92 instance_id: e.instance_id.clone(),
93 token_id: t.id.clone(),
94 node_id: t.node_id.clone(),
95 }));
96 }
97 instance.tokens.extend(new_tokens);
98 }
99 NodeType::UserTask => {
100 instance.tokens[token_idx].status = TokenStatus::Waiting;
101 }
102 NodeType::ExternalTask {
103 task_type,
104 retries,
105 timeout_secs,
106 } => {
107 if let Some(ref ext_store) = ctx.external_task_store {
108 let variables = instance.variables.clone();
109 let _ = ext_store
110 .create(
111 &e.token_id,
112 &e.instance_id,
113 task_type,
114 *retries,
115 *timeout_secs,
116 variables,
117 )
118 .await;
119 }
120 instance.tokens[token_idx].status = TokenStatus::Waiting;
121 }
122 NodeType::ExclusiveGateway => {
123 instance.tokens[token_idx].status = TokenStatus::Waiting;
124 if let Some(t) = evaluate_exclusive_gateway(node, &instance.variables) {
125 out.push(EngineEvent::TokenArrived(payloads::TokenArrived {
126 instance_id: e.instance_id.clone(),
127 token_id: t.id.clone(),
128 node_id: t.node_id.clone(),
129 }));
130 instance.tokens.push(t);
131 }
132 }
133 NodeType::End => {
134 instance.state = InstanceState::Completed;
135 instance.tokens.remove(token_idx);
136 let _ = process_store.save(&instance).await;
137 return vec![EngineEvent::ProcessCompleted(payloads::ProcessCompleted {
138 instance_id: e.instance_id.clone(),
139 })];
140 }
141 NodeType::ParallelFork => {
142 let group_id = uuid::Uuid::new_v4().to_string();
143 if let Some(ref join_repo) = ctx.parallel_join_repo {
144 let expected = node.outgoing_edges.len() as u32;
145 let _ = join_repo.ensure_group(&group_id, expected).await;
146 }
147 instance.tokens[token_idx].status = TokenStatus::Waiting;
148 let new_tokens = move_token_with_group(node, group_id.clone());
149 for t in &new_tokens {
150 out.push(EngineEvent::TokenArrived(payloads::TokenArrived {
151 instance_id: e.instance_id.clone(),
152 token_id: t.id.clone(),
153 node_id: t.node_id.clone(),
154 }));
155 }
156 instance.tokens.extend(new_tokens);
157 }
158 NodeType::ParallelJoin { expected } => {
159 let group_id = instance.tokens[token_idx]
160 .parallel_group_id
161 .clone()
162 .unwrap_or_default();
163 let done = if let Some(ref join_repo) = ctx.parallel_join_repo {
164 join_repo.try_join(&group_id).await.unwrap_or(false)
165 } else {
166 let key = format!("{}:{}:{}", e.instance_id, e.node_id, group_id);
167 let mut state = self.join_state.lock().unwrap();
168 let (exp, arrived) = state
169 .entry(key.clone())
170 .or_insert((*expected, HashSet::new()));
171 arrived.insert(e.token_id.clone());
172 let done = arrived.len() >= *exp;
173 if done {
174 state.remove(&key);
175 }
176 done
177 };
178
179 if done {
180 instance.tokens.retain(|t| {
181 !(t.node_id == e.node_id
182 && t.parallel_group_id.as_deref() == Some(group_id.as_str()))
183 });
184 let new_tokens = move_token(node);
185 for t in &new_tokens {
186 out.push(EngineEvent::TokenArrived(payloads::TokenArrived {
187 instance_id: e.instance_id.clone(),
188 token_id: t.id.clone(),
189 node_id: t.node_id.clone(),
190 }));
191 }
192 instance.tokens.extend(new_tokens);
193 } else {
194 instance.tokens[token_idx].status = TokenStatus::Waiting;
195 }
196 }
197 }
198
199 let _ = process_store.save(&instance).await;
200 out
201 }
202}