aion_worker/
activity.rs

1//! `Activity` trait, `ActivityFailure`, and typed registration.
2
3use std::any::Any;
4use std::collections::{BTreeMap, BTreeSet};
5use std::future::Future;
6use std::panic::AssertUnwindSafe;
7use std::pin::Pin;
8
9use aion_core::{ActivityError, ActivityErrorKind, Payload};
10use async_trait::async_trait;
11use futures::FutureExt;
12use serde::Serialize;
13use serde::de::DeserializeOwned;
14use tracing::error;
15
16use crate::context::ActivityContext;
17use crate::error::{MissingActivityHandler, WorkerError};
18use crate::protocol::ActivityTask;
19use crate::runtime::loop_::{ActivityDispatcher, DispatchOutcome};
20
21/// Explicit retryability classification for an activity failure.
22#[derive(Clone, Debug, PartialEq, Eq)]
23pub enum Classification {
24    /// The engine may retry the activity according to policy.
25    Retryable,
26    /// The activity failure is permanent and must not be retried.
27    Terminal,
28}
29
30/// Handler-returned failure with explicit retryability classification.
31#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
32#[error("{message}")]
33pub struct ActivityFailure {
34    classification: Classification,
35    message: String,
36    detail: Option<Payload>,
37}
38
39impl ActivityFailure {
40    /// Creates a retryable activity failure.
41    #[must_use]
42    pub fn retryable(message: impl Into<String>) -> Self {
43        Self::new(Classification::Retryable, message, None)
44    }
45
46    /// Creates a terminal activity failure.
47    #[must_use]
48    pub fn terminal(message: impl Into<String>) -> Self {
49        Self::new(Classification::Terminal, message, None)
50    }
51
52    /// Attaches opaque structured detail to this failure.
53    #[must_use]
54    pub fn with_detail(mut self, detail: Payload) -> Self {
55        self.detail = Some(detail);
56        self
57    }
58
59    /// Returns the explicit retryability classification.
60    #[must_use]
61    pub const fn classification(&self) -> &Classification {
62        &self.classification
63    }
64
65    /// Returns the human-readable failure message.
66    #[must_use]
67    pub fn message(&self) -> &str {
68        &self.message
69    }
70
71    /// Returns the optional structured failure detail.
72    #[must_use]
73    pub const fn detail(&self) -> Option<&Payload> {
74        self.detail.as_ref()
75    }
76
77    fn new(
78        classification: Classification,
79        message: impl Into<String>,
80        detail: Option<Payload>,
81    ) -> Self {
82        Self {
83            classification,
84            message: message.into(),
85            detail,
86        }
87    }
88}
89
90impl From<Classification> for ActivityErrorKind {
91    fn from(value: Classification) -> Self {
92        match value {
93            Classification::Retryable => Self::Retryable,
94            Classification::Terminal => Self::Terminal,
95        }
96    }
97}
98
99impl From<ActivityFailure> for ActivityError {
100    fn from(value: ActivityFailure) -> Self {
101        Self {
102            kind: ActivityErrorKind::from(value.classification),
103            message: value.message,
104            details: value.detail,
105        }
106    }
107}
108
109/// Boxed future returned by a typed activity handler.
110pub type HandlerFuture<'context, Output> =
111    Pin<Box<dyn Future<Output = Result<Output, ActivityFailure>> + Send + 'context>>;
112
113type BoxedHandler<Input, Output> = Box<
114    dyn for<'context> Fn(Input, &'context ActivityContext) -> HandlerFuture<'context, Output>
115        + Send
116        + Sync,
117>;
118
119/// Registry of typed activity handlers keyed by activity-type name.
120#[derive(Default)]
121pub struct ActivityRegistry {
122    handlers: BTreeMap<String, Box<dyn ErasedActivityHandler>>,
123    descriptors: BTreeMap<String, aion_package::ActivityDescriptor>,
124}
125
126impl ActivityRegistry {
127    /// Creates an empty activity registry.
128    #[must_use]
129    pub fn new() -> Self {
130        Self::default()
131    }
132
133    /// Registers one typed activity handler under an activity-type name.
134    ///
135    /// # Errors
136    ///
137    /// Returns [`WorkerError::Registration`] when the name is already registered.
138    pub fn register_activity<Input, Output, Handler>(
139        mut self,
140        activity_type: impl Into<String>,
141        handler: Handler,
142    ) -> Result<Self, WorkerError>
143    where
144        Input: Serialize + DeserializeOwned + Send + Sync + 'static,
145        Output: Serialize + Send + Sync + 'static,
146        Handler: for<'context> Fn(Input, &'context ActivityContext) -> HandlerFuture<'context, Output>
147            + Send
148            + Sync
149            + 'static,
150    {
151        let activity_type = activity_type.into();
152        if self.handlers.contains_key(&activity_type) {
153            return Err(WorkerError::registration(DuplicateActivityType {
154                activity_type,
155            }));
156        }
157        self.handlers
158            .insert(activity_type, Box::new(TypedHandler::new(handler)));
159        Ok(self)
160    }
161
162    /// Registers a typed handler and mechanically derives its wire contract
163    /// from the same concrete serde types used at dispatch.
164    ///
165    /// # Errors
166    ///
167    /// Returns [`WorkerError::Registration`] for a duplicate name or
168    /// [`WorkerError::Encode`] if a generated schema cannot be represented as JSON.
169    pub fn register_activity_with_contract<Input, Output, Handler>(
170        mut self,
171        activity_type: impl Into<String>,
172        handler: Handler,
173    ) -> Result<Self, WorkerError>
174    where
175        Input: Serialize + DeserializeOwned + schemars::JsonSchema + Send + Sync + 'static,
176        Output: Serialize + schemars::JsonSchema + Send + Sync + 'static,
177        Handler: for<'context> Fn(Input, &'context ActivityContext) -> HandlerFuture<'context, Output>
178            + Send
179            + Sync
180            + 'static,
181    {
182        let activity_type = activity_type.into();
183        let descriptor = activity_descriptor::<Input, Output>(activity_type.clone())?;
184        self = self.register_activity(activity_type.clone(), handler)?;
185        self.descriptors.insert(activity_type, descriptor);
186        Ok(self)
187    }
188
189    /// Registers a typed handler together with an EXPLICIT wire descriptor.
190    ///
191    /// For an activity whose schemas are owned by a declaration rather than by
192    /// Rust types, [`Self::register_activity_with_contract`] cannot help: there
193    /// is no `schemars` type to derive from, because the types live in the
194    /// `.awl` document. The handler is therefore untyped (`serde_json::Value`
195    /// in and out) while the advertisement is exact, supplied by whoever read
196    /// the declaration.
197    ///
198    /// Registering a handler with no descriptor is what makes a worker invisible
199    /// to contract admission — it advertises nothing, and a queue carrying any
200    /// deployed contract refuses it. This is the path that keeps a
201    /// declaration-driven worker admissible.
202    ///
203    /// # Errors
204    ///
205    /// Returns [`WorkerError::Registration`] when the name is already
206    /// registered, or when `descriptor` names a different activity than
207    /// `activity_type` — a descriptor advertised under the wrong name would
208    /// promise one action's schema for another's handler.
209    pub fn register_activity_with_descriptor<Input, Output, Handler>(
210        mut self,
211        activity_type: impl Into<String>,
212        descriptor: aion_package::ActivityDescriptor,
213        handler: Handler,
214    ) -> Result<Self, WorkerError>
215    where
216        Input: Serialize + DeserializeOwned + Send + Sync + 'static,
217        Output: Serialize + Send + Sync + 'static,
218        Handler: for<'context> Fn(Input, &'context ActivityContext) -> HandlerFuture<'context, Output>
219            + Send
220            + Sync
221            + 'static,
222    {
223        let activity_type = activity_type.into();
224        if descriptor.name != activity_type {
225            return Err(WorkerError::registration(DescriptorNameMismatch {
226                activity_type,
227                descriptor: descriptor.name,
228            }));
229        }
230        self = self.register_activity(activity_type.clone(), handler)?;
231        self.descriptors.insert(activity_type, descriptor);
232        Ok(self)
233    }
234
235    /// Returns true when no activity handlers have been registered.
236    #[must_use]
237    pub fn is_empty(&self) -> bool {
238        self.handlers.is_empty()
239    }
240
241    /// Returns the registered activity-type names in deterministic order.
242    #[must_use]
243    pub fn activity_types(&self) -> BTreeSet<String> {
244        self.handlers.keys().cloned().collect()
245    }
246
247    /// Returns committed typed activity descriptors in deterministic name order.
248    #[must_use]
249    pub fn activity_descriptors(&self) -> Vec<aion_package::ActivityDescriptor> {
250        self.descriptors.values().cloned().collect()
251    }
252}
253
254/// Mechanically derives a wire descriptor from concrete serde activity types.
255///
256/// # Errors
257///
258/// Returns [`WorkerError::Encode`] if a generated schema cannot be represented
259/// as JSON.
260pub fn activity_descriptor<Input, Output>(
261    name: impl Into<String>,
262) -> Result<aion_package::ActivityDescriptor, WorkerError>
263where
264    Input: schemars::JsonSchema,
265    Output: schemars::JsonSchema,
266{
267    let schema = || {
268        schemars::generate::SchemaSettings::draft2020_12()
269            .with(|settings| {
270                // AWL contracts carry their dialect at the document root; the
271                // registration envelope already fixes that dialect, so draft
272                // metadata and presentation-only `$defs` must not create drift.
273                settings.meta_schema = None;
274                settings.inline_subschemas = true;
275            })
276            .into_generator()
277    };
278    Ok(aion_package::ActivityDescriptor {
279        name: name.into(),
280        input_schema: serde_json::to_value(schema().into_root_schema_for::<Input>())
281            .map_err(WorkerError::encode)?,
282        output_schema: serde_json::to_value(schema().into_root_schema_for::<Output>())
283            .map_err(WorkerError::encode)?,
284    })
285}
286
287#[async_trait]
288impl ActivityDispatcher for ActivityRegistry {
289    async fn dispatch(
290        &self,
291        task: ActivityTask,
292        context: ActivityContext,
293    ) -> Result<DispatchOutcome, WorkerError> {
294        let Some(handler) = self.handlers.get(&task.activity_type) else {
295            return Err(WorkerError::registration(MissingActivityHandler {
296                activity_type: task.activity_type,
297            }));
298        };
299        handler.dispatch(task, context).await
300    }
301
302    fn activity_types(&self) -> BTreeSet<String> {
303        self.activity_types()
304    }
305}
306
307/// Backwards-compatible name for the typed activity registry used by the runtime.
308pub type TypedActivityDispatcher = ActivityRegistry;
309
310/// Decodes a payload into a typed value using the payload content-type tag.
311///
312/// # Errors
313///
314/// Returns [`WorkerError::Decode`] when the payload tag or bytes cannot produce
315/// the requested type.
316pub fn decode_payload<T>(payload: &Payload) -> Result<T, WorkerError>
317where
318    T: DeserializeOwned,
319{
320    let value = payload.to_json().map_err(WorkerError::decode)?;
321    serde_json::from_value(value).map_err(WorkerError::decode)
322}
323
324/// Encodes a typed value into the baseline JSON payload codec.
325///
326/// # Errors
327///
328/// Returns [`WorkerError::Encode`] when the value cannot be serialized.
329pub fn encode_payload<T>(value: &T) -> Result<Payload, WorkerError>
330where
331    T: Serialize,
332{
333    let value = serde_json::to_value(value).map_err(WorkerError::encode)?;
334    Payload::from_json(&value).map_err(WorkerError::encode)
335}
336
337#[async_trait]
338trait ErasedActivityHandler: Send + Sync {
339    async fn dispatch(
340        &self,
341        task: ActivityTask,
342        context: ActivityContext,
343    ) -> Result<DispatchOutcome, WorkerError>;
344}
345
346struct TypedHandler<Input, Output> {
347    handler: BoxedHandler<Input, Output>,
348}
349
350impl<Input, Output> TypedHandler<Input, Output> {
351    fn new(
352        handler: impl for<'context> Fn(
353            Input,
354            &'context ActivityContext,
355        ) -> HandlerFuture<'context, Output>
356        + Send
357        + Sync
358        + 'static,
359    ) -> Self {
360        Self {
361            handler: Box::new(handler),
362        }
363    }
364}
365
366#[async_trait]
367impl<Input, Output> ErasedActivityHandler for TypedHandler<Input, Output>
368where
369    Input: DeserializeOwned + Send + Sync + 'static,
370    Output: Serialize + Send + Sync + 'static,
371{
372    async fn dispatch(
373        &self,
374        task: ActivityTask,
375        context: ActivityContext,
376    ) -> Result<DispatchOutcome, WorkerError> {
377        let input = match decode_payload::<Input>(&task.input) {
378            Ok(input) => input,
379            Err(error) => {
380                error!(
381                    activity_type = %task.activity_type,
382                    activity_id = task.activity_id.sequence_position(),
383                    attempt = task.attempt,
384                    error = %error,
385                    "failed to decode activity input; reporting terminal activity failure"
386                );
387                let failure =
388                    ActivityFailure::terminal(format!("failed to decode activity input: {error}"));
389                return Ok(DispatchOutcome::Failed {
390                    failure: ActivityError::from(failure),
391                });
392            }
393        };
394        let handler_future =
395            match std::panic::catch_unwind(AssertUnwindSafe(|| (self.handler)(input, &context))) {
396                Ok(handler_future) => handler_future,
397                Err(panic) => return Ok(panic_failure(&task, &panic)),
398            };
399        let handler_result = AssertUnwindSafe(handler_future).catch_unwind().await;
400        let outcome = match handler_result {
401            Ok(Ok(output)) => DispatchOutcome::Completed {
402                output: encode_payload(&output)?,
403            },
404            Ok(Err(failure)) => DispatchOutcome::Failed {
405                failure: ActivityError::from(failure),
406            },
407            Err(panic) => panic_failure(&task, &panic),
408        };
409        Ok(outcome)
410    }
411}
412
413fn panic_failure(task: &ActivityTask, panic: &Box<dyn Any + Send>) -> DispatchOutcome {
414    let message = panic_message(panic);
415    error!(
416        activity_type = %task.activity_type,
417        activity_id = task.activity_id.sequence_position(),
418        attempt = task.attempt,
419        panic = %message,
420        "activity handler panicked; reporting retryable activity failure"
421    );
422    DispatchOutcome::Failed {
423        failure: ActivityError::from(ActivityFailure::retryable(format!(
424            "activity handler panicked: {message}"
425        ))),
426    }
427}
428
429fn panic_message(panic: &Box<dyn Any + Send>) -> String {
430    if let Some(message) = panic.downcast_ref::<&str>() {
431        return (*message).to_owned();
432    }
433    if let Some(message) = panic.downcast_ref::<String>() {
434        return message.clone();
435    }
436    String::from("unknown panic payload")
437}
438
439/// Error returned when an activity type is registered more than once.
440#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
441#[error("activity type `{activity_type}` already has a registered handler")]
442pub struct DuplicateActivityType {
443    /// Duplicate activity type name.
444    pub activity_type: String,
445}
446
447/// Error returned when an explicit descriptor names a different activity than
448/// the one it is registered under.
449///
450/// Advertising a schema under the wrong name promises one action's shape for
451/// another action's handler, which contract admission would then accept — so
452/// the mismatch is refused at registration instead.
453#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
454#[error(
455    "activity type `{activity_type}` was registered with a descriptor naming `{descriptor}`; a descriptor must name the activity it describes"
456)]
457pub struct DescriptorNameMismatch {
458    /// The name the handler is registered under.
459    pub activity_type: String,
460    /// The name the supplied descriptor carries.
461    pub descriptor: String,
462}
463
464#[cfg(test)]
465mod tests {
466    use aion_core::{ActivityError, ActivityId, ContentType, WorkflowId};
467    use aion_proto::{
468        ProtoActivityError, ProtoActivityErrorKind, ProtoActivityId, ProtoActivityTask,
469        ProtoPayload, ProtoWorkflowId,
470    };
471    use serde::{Deserialize, Serialize};
472
473    use super::{ActivityFailure, ActivityRegistry, decode_payload, encode_payload};
474    use crate::WorkerError;
475    use crate::runtime::{ActivityDispatcher, DispatchOutcome};
476
477    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
478    struct TestInput {
479        value: i32,
480    }
481
482    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
483    struct TestOutput {
484        doubled: i32,
485    }
486
487    #[test]
488    fn retryable_and_terminal_failures_map_to_distinct_wire_classifications() {
489        let retryable = ActivityFailure::retryable("temporary outage");
490        let terminal = ActivityFailure::terminal("invalid request");
491
492        let retryable_core = ActivityError::from(retryable);
493        let terminal_core = ActivityError::from(terminal);
494        let retryable_wire = ProtoActivityError::from(retryable_core);
495        let terminal_wire = ProtoActivityError::from(terminal_core);
496
497        assert_eq!(
498            retryable_wire.kind,
499            ProtoActivityErrorKind::Retryable as i32
500        );
501        assert_eq!(terminal_wire.kind, ProtoActivityErrorKind::Terminal as i32);
502    }
503
504    #[tokio::test]
505    async fn typed_activity_round_trips_through_registry() -> Result<(), WorkerError> {
506        let registry =
507            ActivityRegistry::new().register_activity("double", |input: TestInput, context| {
508                Box::pin(async move {
509                    assert_eq!(context.attempt(), 1);
510                    Ok(TestOutput {
511                        doubled: input.value * 2,
512                    })
513                })
514            })?;
515        let task = proto_task("double", &TestInput { value: 21 })?;
516        let (context, cancellation) = crate::ActivityContext::for_workflow(
517            Some(WorkflowId::new_v4()),
518            ActivityId::from_sequence_position(99),
519            1,
520            None,
521        );
522        drop(cancellation);
523
524        let outcome = registry.dispatch(task.try_into()?, context).await?;
525
526        let DispatchOutcome::Completed { output } = outcome else {
527            return Err(WorkerError::decode(UnexpectedFailure));
528        };
529        assert_eq!(output.content_type(), &ContentType::Json);
530        let decoded: TestOutput = decode_payload(&output)?;
531        assert_eq!(decoded, TestOutput { doubled: 42 });
532        Ok(())
533    }
534
535    #[test]
536    fn duplicate_activity_registration_is_rejected() -> Result<(), WorkerError> {
537        let registry =
538            ActivityRegistry::new().register_activity("double", |input: TestInput, context| {
539                Box::pin(async move {
540                    let _ = context;
541                    Ok(TestOutput {
542                        doubled: input.value * 2,
543                    })
544                })
545            })?;
546
547        let error = registry
548            .register_activity("double", |input: TestInput, context| {
549                Box::pin(async move {
550                    let _ = context;
551                    Ok(TestOutput {
552                        doubled: input.value,
553                    })
554                })
555            })
556            .err()
557            .ok_or_else(|| WorkerError::decode(UnexpectedFailure))?;
558
559        assert!(
560            error
561                .to_string()
562                .contains("already has a registered handler")
563        );
564        Ok(())
565    }
566
567    fn proto_task(
568        activity_type: &str,
569        input: &TestInput,
570    ) -> Result<ProtoActivityTask, WorkerError> {
571        Ok(ProtoActivityTask {
572            workflow_id: Some(ProtoWorkflowId::from(WorkflowId::new_v4())),
573            activity_id: Some(ProtoActivityId::from(ActivityId::from_sequence_position(1))),
574            run_id: Some(aion_proto::ProtoRunId::from(aion_core::RunId::new_v4())),
575            activity_type: activity_type.to_owned(),
576            input: Some(ProtoPayload::from(encode_payload(&input)?)),
577            attempt: 1,
578            completion_token: String::from("generation-1"),
579            idempotency_key: String::from("effect-key"),
580            labels: std::collections::HashMap::new(),
581        })
582    }
583
584    #[derive(Debug, thiserror::Error)]
585    #[error("expected completed activity outcome")]
586    struct UnexpectedFailure;
587}
aion_worker/activity.rs

aion_worker/
activity.rs
