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    /// Revoke `token` if it is still current, without disturbing a newer retry.
116    ///
117    /// # Errors
118    ///
119    /// Returns [`ServerError::LockPoisoned`] when fence state cannot be trusted.
120    pub fn revoke(
121        &self,
122        workflow_id: &WorkflowId,
123        activity_id: &ActivityId,
124        token: &CompletionToken,
125    ) -> Result<(), ServerError> {
126        let key = (workflow_id.clone(), activity_id.clone());
127        let mut state = self.state()?;
128        if state.get(&key) == Some(token) {
129            state.remove(&key);
130        }
131        Ok(())
132    }
133
134    /// Revoke whichever generation is current while parking for recovery.
135    ///
136    /// # Errors
137    ///
138    /// Returns [`ServerError::LockPoisoned`] when fence state cannot be trusted.
139    pub fn revoke_current(
140        &self,
141        workflow_id: &WorkflowId,
142        activity_id: &ActivityId,
143    ) -> Result<(), ServerError> {
144        self.state()?
145            .remove(&(workflow_id.clone(), activity_id.clone()));
146        Ok(())
147    }
148
149    fn state(&self) -> Result<MutexGuard<'_, HashMap<ExecutionKey, CompletionToken>>, ServerError> {
150        self.current
151            .lock()
152            .map_err(|_| ServerError::lock_poisoned("activity completion fences"))
153    }
154}
155
156/// Derive the stable external-effect key for one action site in one workflow run.
157///
158/// Attempts and execution generations are intentionally absent. The domain tag,
159/// workflow id, run id, and activity ordinal are length-unambiguous fixed-width
160/// inputs to SHA-256.
161#[must_use]
162pub fn idempotency_key(
163    workflow_id: &WorkflowId,
164    run_id: &RunId,
165    activity_id: &ActivityId,
166) -> String {
167    let mut hasher = Sha256::new();
168    hasher.update(b"aion.activity.idempotency.v1\0");
169    hasher.update(workflow_id.as_uuid().as_bytes());
170    hasher.update(run_id.as_uuid().as_bytes());
171    hasher.update(activity_id.sequence_position().to_be_bytes());
172    encode_hex(&hasher.finalize())
173}
174
175fn encode_hex(bytes: &[u8]) -> String {
176    const DIGITS: &[u8; 16] = b"0123456789abcdef";
177    let mut encoded = String::with_capacity(bytes.len() * 2);
178    for byte in bytes {
179        encoded.push(char::from(DIGITS[usize::from(byte >> 4)]));
180        encoded.push(char::from(DIGITS[usize::from(byte & 0x0f)]));
181    }
182    encoded
183}
184
185fn rejection(
186    workflow_id: &WorkflowId,
187    activity_id: &ActivityId,
188    reason: CompletionRejectionReason,
189) -> ServerError {
190    ServerError::ActivityCompletionRejected {
191        workflow_id: workflow_id.clone(),
192        activity_id: activity_id.clone(),
193        reason,
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use super::{CompletionFences, CompletionToken, idempotency_key};
200    use crate::error::{CompletionRejectionReason, ServerError};
201    use aion_core::{ActivityId, RunId, WorkflowId};
202
203    type TestResult = Result<(), Box<dyn std::error::Error>>;
204
205    #[test]
206    fn idempotency_key_is_attempt_independent_and_site_run_scoped() {
207        let workflow = WorkflowId::new_v4();
208        let run_a = RunId::new_v4();
209        let run_b = RunId::new_v4();
210        let site_a = ActivityId::from_sequence_position(7);
211        let site_b = ActivityId::from_sequence_position(8);
212
213        let first_attempt = idempotency_key(&workflow, &run_a, &site_a);
214        let fifth_attempt = idempotency_key(&workflow, &run_a, &site_a);
215        assert_eq!(first_attempt, fifth_attempt);
216        assert_ne!(first_attempt, idempotency_key(&workflow, &run_a, &site_b));
217        assert_ne!(first_attempt, idempotency_key(&workflow, &run_b, &site_a));
218    }
219
220    #[test]
221    fn issuing_a_retry_rejects_the_stale_generation() -> TestResult {
222        let fences = CompletionFences::default();
223        let workflow = WorkflowId::new_v4();
224        let activity = ActivityId::from_sequence_position(3);
225        let stale = fences.issue(&workflow, &activity)?;
226        let current = fences.issue(&workflow, &activity)?;
227
228        let rejected = fences.accept(&workflow, &activity, &stale);
229        assert!(matches!(
230            rejected,
231            Err(ServerError::ActivityCompletionRejected {
232                reason: CompletionRejectionReason::StaleGeneration,
233                ..
234            })
235        ));
236        fences.accept(&workflow, &activity, &current)?;
237        Ok(())
238    }
239
240    #[test]
241    fn accepted_generation_is_consumed_exactly_once() -> TestResult {
242        let fences = CompletionFences::default();
243        let workflow = WorkflowId::new_v4();
244        let activity = ActivityId::from_sequence_position(4);
245        let token = fences.issue(&workflow, &activity)?;
246
247        fences.accept(&workflow, &activity, &token)?;
248        let duplicate = fences.accept(&workflow, &activity, &token);
249        assert!(matches!(
250            duplicate,
251            Err(ServerError::ActivityCompletionRejected {
252                reason: CompletionRejectionReason::NoCurrentGeneration,
253                ..
254            })
255        ));
256        Ok(())
257    }
258
259    #[test]
260    fn revoking_an_old_generation_does_not_remove_its_replacement() -> TestResult {
261        let fences = CompletionFences::default();
262        let workflow = WorkflowId::new_v4();
263        let activity = ActivityId::from_sequence_position(6);
264        let old = fences.issue(&workflow, &activity)?;
265        let replacement = fences.issue(&workflow, &activity)?;
266
267        fences.revoke(&workflow, &activity, &old)?;
268        fences.accept(&workflow, &activity, &replacement)?;
269        Ok(())
270    }
271
272    #[test]
273    fn an_empty_wire_token_is_a_typed_compatibility_refusal() {
274        let workflow = WorkflowId::new_v4();
275        let activity = ActivityId::from_sequence_position(9);
276        let rejected = CompletionToken::from_wire(&workflow, &activity, String::new());
277        assert!(matches!(
278            rejected,
279            Err(ServerError::ActivityCompletionRejected {
280                reason: CompletionRejectionReason::MissingCompletionToken,
281                ..
282            })
283        ));
284    }
285
286    #[test]
287    fn a_pre_recovery_generation_is_rejected_after_recovery() -> TestResult {
288        let before_recovery = CompletionFences::default();
289        let workflow = WorkflowId::new_v4();
290        let activity = ActivityId::from_sequence_position(5);
291        let stale = before_recovery.issue(&workflow, &activity)?;
292
293        let after_recovery = CompletionFences::default();
294        let current = after_recovery.issue(&workflow, &activity)?;
295        let rejected = after_recovery.accept(&workflow, &activity, &stale);
296        assert!(matches!(
297            rejected,
298            Err(ServerError::ActivityCompletionRejected {
299                reason: CompletionRejectionReason::StaleGeneration,
300                ..
301            })
302        ));
303        after_recovery.accept(&workflow, &activity, &current)?;
304        Ok(())
305    }
306}
aion_server/worker/envelope.rs

aion_server/worker/
envelope.rs
