aion_server/worker/
envelope.rs

1//! Activity task-envelope generation and completion fencing.
2
3use std::collections::HashMap;
4use std::sync::{Arc, Mutex, MutexGuard};
5
6use aion_core::{ActivityId, RunId, WorkflowId};
7use sha2::{Digest, Sha256};
8use uuid::Uuid;
9
10use crate::error::{CompletionRejectionReason, ServerError};
11
12type ExecutionKey = (WorkflowId, ActivityId);
13
14/// Opaque proof that a worker owns one dispatched execution generation.
15#[derive(Clone, Debug, Eq, PartialEq)]
16pub struct CompletionToken(String);
17
18impl CompletionToken {
19    /// Parse a worker-echoed token without assigning meaning to its contents.
20    ///
21    /// # Errors
22    ///
23    /// Returns a typed compatibility refusal when a pre-fencing worker omits it.
24    pub fn from_wire(
25        workflow_id: &WorkflowId,
26        activity_id: &ActivityId,
27        value: String,
28    ) -> Result<Self, ServerError> {
29        if value.is_empty() {
30            return Err(rejection(
31                workflow_id,
32                activity_id,
33                CompletionRejectionReason::MissingCompletionToken,
34            ));
35        }
36        Ok(Self(value))
37    }
38
39    /// Return the opaque wire representation.
40    #[must_use]
41    pub fn as_str(&self) -> &str {
42        &self.0
43    }
44
45    /// Build a non-wire token for crate-local unit fixtures.
46    #[cfg(test)]
47    #[must_use]
48    pub(crate) fn for_test() -> Self {
49        Self("test-generation".to_owned())
50    }
51}
52
53/// Process-incarnation registry of the only generations allowed to complete.
54///
55/// Recovery deliberately starts with an empty registry. Re-dispatch issues a
56/// fresh token, so a result carrying any pre-recovery token is refused whether
57/// it arrives before or after the new dispatch.
58#[derive(Clone, Debug, Default)]
59pub struct CompletionFences {
60    current: Arc<Mutex<HashMap<ExecutionKey, CompletionToken>>>,
61}
62
63impl CompletionFences {
64    /// Supersede any prior generation and return the newly authorized token.
65    ///
66    /// # Errors
67    ///
68    /// Returns [`ServerError::LockPoisoned`] when fence state cannot be trusted.
69    pub fn issue(
70        &self,
71        workflow_id: &WorkflowId,
72        activity_id: &ActivityId,
73    ) -> Result<CompletionToken, ServerError> {
74        let token = CompletionToken(Uuid::new_v4().to_string());
75        self.state()?
76            .insert((workflow_id.clone(), activity_id.clone()), token.clone());
77        Ok(token)
78    }
79
80    /// Consume the current generation only when `submitted` exactly matches it.
81    ///
82    /// Consumption and comparison share one mutex critical section, making two
83    /// concurrent submissions unable to both become truth.
84    ///
85    /// # Errors
86    ///
87    /// Returns a typed rejection for a missing/current-generation mismatch, or
88    /// [`ServerError::LockPoisoned`] when fence state cannot be trusted.
89    pub fn accept(
90        &self,
91        workflow_id: &WorkflowId,
92        activity_id: &ActivityId,
93        submitted: &CompletionToken,
94    ) -> Result<(), ServerError> {
95        let key = (workflow_id.clone(), activity_id.clone());
96        let mut state = self.state()?;
97        let Some(current) = state.get(&key) else {
98            return Err(rejection(
99                workflow_id,
100                activity_id,
101                CompletionRejectionReason::NoCurrentGeneration,
102            ));
103        };
104        if current != submitted {
105            return Err(rejection(
106                workflow_id,
107                activity_id,
108                CompletionRejectionReason::StaleGeneration,
109            ));
110        }
111        state.remove(&key);
112        Ok(())
113    }
114
115    /// Restore a consumed generation after accepted-path settlement fails.
116    ///
117    /// A concurrently issued newer generation always wins; the accepted token is
118    /// restored only while the execution site has no current generation.
119    ///
120    /// # Errors
121    ///
122    /// Returns [`ServerError::LockPoisoned`] when fence state cannot be trusted.
123    pub fn restore_if_absent(
124        &self,
125        workflow_id: &WorkflowId,
126        activity_id: &ActivityId,
127        token: &CompletionToken,
128    ) -> Result<(), ServerError> {
129        self.state()?
130            .entry((workflow_id.clone(), activity_id.clone()))
131            .or_insert_with(|| token.clone());
132        Ok(())
133    }
134
135    /// Revoke `token` if it is still current, without disturbing a newer retry.
136    ///
137    /// # Errors
138    ///
139    /// Returns [`ServerError::LockPoisoned`] when fence state cannot be trusted.
140    pub fn revoke(
141        &self,
142        workflow_id: &WorkflowId,
143        activity_id: &ActivityId,
144        token: &CompletionToken,
145    ) -> Result<(), ServerError> {
146        let key = (workflow_id.clone(), activity_id.clone());
147        let mut state = self.state()?;
148        if state.get(&key) == Some(token) {
149            state.remove(&key);
150        }
151        Ok(())
152    }
153
154    /// Revoke whichever generation is current while parking for recovery.
155    ///
156    /// # Errors
157    ///
158    /// Returns [`ServerError::LockPoisoned`] when fence state cannot be trusted.
159    pub fn revoke_current(
160        &self,
161        workflow_id: &WorkflowId,
162        activity_id: &ActivityId,
163    ) -> Result<(), ServerError> {
164        self.state()?
165            .remove(&(workflow_id.clone(), activity_id.clone()));
166        Ok(())
167    }
168
169    fn state(&self) -> Result<MutexGuard<'_, HashMap<ExecutionKey, CompletionToken>>, ServerError> {
170        self.current
171            .lock()
172            .map_err(|_| ServerError::lock_poisoned("activity completion fences"))
173    }
174}
175
176/// Derive the stable external-effect key for one action site in one workflow run.
177///
178/// Attempts and execution generations are intentionally absent. The domain tag,
179/// workflow id, run id, and activity ordinal are length-unambiguous fixed-width
180/// inputs to SHA-256.
181#[must_use]
182pub fn idempotency_key(
183    workflow_id: &WorkflowId,
184    run_id: &RunId,
185    activity_id: &ActivityId,
186) -> String {
187    let mut hasher = Sha256::new();
188    hasher.update(b"aion.activity.idempotency.v1\0");
189    hasher.update(workflow_id.as_uuid().as_bytes());
190    hasher.update(run_id.as_uuid().as_bytes());
191    hasher.update(activity_id.sequence_position().to_be_bytes());
192    encode_hex(&hasher.finalize())
193}
194
195fn encode_hex(bytes: &[u8]) -> String {
196    const DIGITS: &[u8; 16] = b"0123456789abcdef";
197    let mut encoded = String::with_capacity(bytes.len() * 2);
198    for byte in bytes {
199        encoded.push(char::from(DIGITS[usize::from(byte >> 4)]));
200        encoded.push(char::from(DIGITS[usize::from(byte & 0x0f)]));
201    }
202    encoded
203}
204
205fn rejection(
206    workflow_id: &WorkflowId,
207    activity_id: &ActivityId,
208    reason: CompletionRejectionReason,
209) -> ServerError {
210    ServerError::ActivityCompletionRejected {
211        workflow_id: workflow_id.clone(),
212        activity_id: activity_id.clone(),
213        reason,
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::{CompletionFences, CompletionToken, idempotency_key};
220    use crate::error::{CompletionRejectionReason, ServerError};
221    use aion_core::{ActivityId, RunId, WorkflowId};
222
223    type TestResult = Result<(), Box<dyn std::error::Error>>;
224
225    #[test]
226    fn idempotency_key_is_attempt_independent_and_site_run_scoped() {
227        let workflow = WorkflowId::new_v4();
228        let run_a = RunId::new_v4();
229        let run_b = RunId::new_v4();
230        let site_a = ActivityId::from_sequence_position(7);
231        let site_b = ActivityId::from_sequence_position(8);
232
233        let first_attempt = idempotency_key(&workflow, &run_a, &site_a);
234        let fifth_attempt = idempotency_key(&workflow, &run_a, &site_a);
235        assert_eq!(first_attempt, fifth_attempt);
236        assert_ne!(first_attempt, idempotency_key(&workflow, &run_a, &site_b));
237        assert_ne!(first_attempt, idempotency_key(&workflow, &run_b, &site_a));
238    }
239
240    #[test]
241    fn issuing_a_retry_rejects_the_stale_generation() -> TestResult {
242        let fences = CompletionFences::default();
243        let workflow = WorkflowId::new_v4();
244        let activity = ActivityId::from_sequence_position(3);
245        let stale = fences.issue(&workflow, &activity)?;
246        let current = fences.issue(&workflow, &activity)?;
247
248        let rejected = fences.accept(&workflow, &activity, &stale);
249        assert!(matches!(
250            rejected,
251            Err(ServerError::ActivityCompletionRejected {
252                reason: CompletionRejectionReason::StaleGeneration,
253                ..
254            })
255        ));
256        fences.accept(&workflow, &activity, &current)?;
257        Ok(())
258    }
259
260    #[test]
261    fn accepted_generation_is_consumed_exactly_once() -> TestResult {
262        let fences = CompletionFences::default();
263        let workflow = WorkflowId::new_v4();
264        let activity = ActivityId::from_sequence_position(4);
265        let token = fences.issue(&workflow, &activity)?;
266
267        fences.accept(&workflow, &activity, &token)?;
268        let duplicate = fences.accept(&workflow, &activity, &token);
269        assert!(matches!(
270            duplicate,
271            Err(ServerError::ActivityCompletionRejected {
272                reason: CompletionRejectionReason::NoCurrentGeneration,
273                ..
274            })
275        ));
276        Ok(())
277    }
278
279    #[test]
280    fn revoking_an_old_generation_does_not_remove_its_replacement() -> TestResult {
281        let fences = CompletionFences::default();
282        let workflow = WorkflowId::new_v4();
283        let activity = ActivityId::from_sequence_position(6);
284        let old = fences.issue(&workflow, &activity)?;
285        let replacement = fences.issue(&workflow, &activity)?;
286
287        fences.revoke(&workflow, &activity, &old)?;
288        fences.accept(&workflow, &activity, &replacement)?;
289        Ok(())
290    }
291
292    #[test]
293    fn an_empty_wire_token_is_a_typed_compatibility_refusal() {
294        let workflow = WorkflowId::new_v4();
295        let activity = ActivityId::from_sequence_position(9);
296        let rejected = CompletionToken::from_wire(&workflow, &activity, String::new());
297        assert!(matches!(
298            rejected,
299            Err(ServerError::ActivityCompletionRejected {
300                reason: CompletionRejectionReason::MissingCompletionToken,
301                ..
302            })
303        ));
304    }
305
306    #[test]
307    fn a_pre_recovery_generation_is_rejected_after_recovery() -> TestResult {
308        let before_recovery = CompletionFences::default();
309        let workflow = WorkflowId::new_v4();
310        let activity = ActivityId::from_sequence_position(5);
311        let stale = before_recovery.issue(&workflow, &activity)?;
312
313        let after_recovery = CompletionFences::default();
314        let current = after_recovery.issue(&workflow, &activity)?;
315        let rejected = after_recovery.accept(&workflow, &activity, &stale);
316        assert!(matches!(
317            rejected,
318            Err(ServerError::ActivityCompletionRejected {
319                reason: CompletionRejectionReason::StaleGeneration,
320                ..
321            })
322        ));
323        after_recovery.accept(&workflow, &activity, &current)?;
324        Ok(())
325    }
326}
aion_server/worker/envelope.rs

aion_server/worker/
envelope.rs
