Repository: slawlor/ractor
Pull request #300 (2024-12-17 18:11:13)

Add proper span handling to factories through extension trait
## Context

With #266, automation span propagation was added to `ractor`. This is very helpful in tracing context through actors, such that contextual information can be persisted through a message's lifecycle. 

However, in factories, this is more subtle due to (a) internal queueing and (b) the worker's requirement to interface with the Factory directly (for ping/pong heartbeats). 

## The problem

Consider a message sent to a factory, which is a `Job`. When the job is picked up, the factory's message handler `handle()` will be instrumented with the span included in the message, silently by the framework for you. However should the factory choose to queue the message, then once that iteration of `handle()` is completed, the span context is lost. Then the factory may utilize a future message event to dequeue that previously enqueued message, and will instrument it with the context from the triggering event. 

In production, it has been noticed that not only can this attach the wrong span, to the wrong message, but it also can cause recursive spans since a worker can send a message to a factory (pong to a ping), which can trigger a queue'd item to be processed, therefore attaching Factory.handle() 'ping' -> Worker.handle() 'pong' -> Factory.handle() 'pong' -> Worker.handle() 'message', and up to infinite depths.

This is not only wrong, but dangerous at runtime (incorrect context risk). Therefore this fix is done by attaching the span to the `Job` directly, inside the `JobOptions` which are generally constructed via `Default::default()` in all call sites. Then a new helper-trait is added, which automatically will instrument the span on the correct Worker.handle() call but also can alleviate a lot of the general worker boilerplate for you (see simplified syntax in examples and tests)

## Important

Given the
1. Removal of the `Eq` bound on the `JobOptions` 
2. Addition of a new field into `JobOptions`
3. General change of the API contract for `JobOptions` to be more read-only, with only allowing the user to set the TTL

this change constitutes a major API break. Although most call sites will see a no-op change, we do need to do a major release to include this change.

# This change

In this PR, I'm doing 2 major changes

1. Adding the `tracing::Span` to the `JobOptions` in order to carry the span for the life of the job, until it's processed in the worker. Additionally we change the API for `JobOptions` to be primarily read-only, with exception only for setting the TTL. The rest is automatically handled for users.
2. Adding a `Worker` trait, which auto-implements the `Actor` trait specifically for a Factory. This trait supports much of the boiler-plate of constructing a worker in a factory and instruments the span if it's provided on the workers handler.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/slawlor/ractor/pull/300?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Sean+Lawlor) Report
Attention: Patch coverage is `65.33333%` with `26 lines` in your changes missing coverage. Please review.
> Project coverage is 82.29%. Comparing base [(`6a08fe2`)](https://app.codecov.io/gh/slawlor/ractor/commit/6a08fe2f50d67b788ab1a9fd052c05fc989bc110?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Sean+Lawlor) to head [(`ee17079`)](https://app.codecov.io/gh/slawlor/ractor/commit/ee17079cca553281da98c1cb89343a655af0477f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Sean+Lawlor).

| [Files with missing lines](https://app.codecov.io/gh/slawlor/ractor/pull/300?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Sean+Lawlor) | Patch % | Lines |
|---|---|---|
| [ractor/src/factory/worker.rs](https://app.codecov.io/gh/slawlor/ractor/pull/300?src=pr&el=tree&filepath=ractor%2Fsrc%2Ffactory%2Fworker.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Sean+Lawlor#diff-cmFjdG9yL3NyYy9mYWN0b3J5L3dvcmtlci5ycw==) | 63.33% | [22 Missing :warning: ](https://app.codecov.io/gh/slawlor/ractor/pull/300?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Sean+Lawlor) |
| [ractor/src/factory/job.rs](https://app.codecov.io/gh/slawlor/ractor/pull/300?src=pr&el=tree&filepath=ractor%2Fsrc%2Ffactory%2Fjob.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Sean+Lawlor#diff-cmFjdG9yL3NyYy9mYWN0b3J5L2pvYi5ycw==) | 55.55% | [4 Missing :warning: ](https://app.codecov.io/gh/slawlor/ractor/pull/300?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Sean+Lawlor) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #300      +/-   ##
==========================================
- Coverage   82.42%   82.29%   -0.14%     
==========================================
  Files          60       60              
  Lines       10924    10983      +59     
==========================================
+ Hits         9004     9038      +34     
- Misses       1920     1945      +25     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/slawlor/ractor/pull/300?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Sean+Lawlor).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Sean+Lawlor).


Patch:
diff --git a/ractor/Cargo.toml b/ractor/Cargo.toml
index 694f3b2..8e3d9e1 100644
--- a/ractor/Cargo.toml
+++ b/ractor/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ractor"
-version = "0.13.5"
+version = "0.14.0"
 authors = ["Sean Lawlor", "Evan Au", "Dillon George"]
 description = "A actor framework for Rust"
 documentation = "https://docs.rs/ractor"
diff --git a/ractor/src/factory/job.rs b/ractor/src/factory/job.rs
index 0e3b76b..b5407ac 100644
--- a/ractor/src/factory/job.rs
+++ b/ractor/src/factory/job.rs
@@ -11,6 +11,7 @@ use std::sync::Arc;
 use std::{hash::Hash, time::SystemTime};
 
 use bon::Builder;
+use tracing::Span;
 
 use crate::{concurrency::Duration, Message};
 use crate::{ActorRef, RpcReplyPort};
@@ -40,27 +41,91 @@ pub trait JobKey: Debug + Hash + Send + Sync + Clone + Eq + PartialEq + 'static
 impl<T: Debug + Hash + Send + Sync + Clone + Eq + PartialEq + 'static> JobKey for T {}
 
 /// Represents options for the specified job
-#[derive(Debug, Eq, PartialEq, Clone)]
+#[derive(Debug, PartialEq, Clone)]
 pub struct JobOptions {
     /// Time job was submitted from the client
-    pub submit_time: SystemTime,
+    submit_time: SystemTime,
     /// Time job was processed by the factory
-    pub factory_time: SystemTime,
+    factory_time: SystemTime,
     /// Time job was sent to a worker
-    pub worker_time: SystemTime,
+    worker_time: SystemTime,
     /// Time-to-live for the job
-    pub ttl: Option<Duration>,
+    ttl: Option<Duration>,
+    /// The parent span we want to propagate to the worker.
+    /// Spans don't propagate over the wire in networks
+    span: Option<Span>,
 }
 
-impl Default for JobOptions {
-    fn default() -> Self {
+impl JobOptions {
+    /// Create a new [JobOptions] instance, optionally supplying the ttl for the job
+    ///
+    /// * `ttl` - The Time-to-live specification for this job, which is the maximum amount
+    ///   of time the job can remain in the factory's (or worker's) queue before being expired
+    ///   and discarded
+    ///
+    /// Returns a new [JobOptions] instance.
+    pub fn new(ttl: Option<Duration>) -> Self {
+        let span = {
+            #[cfg(feature = "message_span_propogation")]
+            {
+                Some(Span::current())
+            }
+            #[cfg(not(feature = "message_span_propogation"))]
+            {
+                None
+            }
+        };
         Self {
             submit_time: SystemTime::now(),
             factory_time: SystemTime::now(),
             worker_time: SystemTime::now(),
-            ttl: None,
+            ttl,
+            span,
         }
     }
+
+    /// Retrieve the TTL for this job
+    pub fn ttl(&self) -> Option<Duration> {
+        self.ttl
+    }
+
+    /// Set the TTL for this job
+    pub fn set_ttl(&mut self, ttl: Option<Duration>) {
+        self.ttl = ttl;
+    }
+
+    /// Time the job was submitted to the factory
+    /// (i.e. the time `cast` was called)
+    pub fn submit_time(&self) -> SystemTime {
+        self.submit_time
+    }
+
+    /// Time the job was dispatched to a worker
+    pub fn worker_time(&self) -> SystemTime {
+        self.worker_time
+    }
+
+    /// Time the job was received by the factory and first either dispatched
+    /// or enqueued to the factory's queue
+    pub fn factory_time(&self) -> SystemTime {
+        self.factory_time
+    }
+
+    /// Clone the [Span] and return it which is attached
+    /// to this [JobOptions] instance.
+    pub fn span(&self) -> Option<Span> {
+        self.span.clone()
+    }
+
+    pub(crate) fn take_span(&mut self) -> Option<Span> {
+        self.span.take()
+    }
+}
+
+impl Default for JobOptions {
+    fn default() -> Self {
+        Self::new(None)
+    }
 }
 
 #[cfg(feature = "cluster")]
@@ -86,7 +151,10 @@ impl BytesConvertable for JobOptions {
 
     fn from_bytes(mut data: Vec<u8>) -> Self {
         if data.len() != 16 {
-            Self::default()
+            Self {
+                span: None,
+                ..Default::default()
+            }
         } else {
             let ttl_bytes = data.split_off(8);
 
@@ -100,6 +168,7 @@ impl BytesConvertable for JobOptions {
                 } else {
                     None
                 },
+                span: None,
                 ..Default::default()
             }
         }
diff --git a/ractor/src/factory/mod.rs b/ractor/src/factory/mod.rs
index 66bd3f5..b40fd2b 100644
--- a/ractor/src/factory/mod.rs
+++ b/ractor/src/factory/mod.rs
@@ -66,53 +66,38 @@
 //! /// the business logic for each message that will be done in parallel.
 //! struct ExampleWorker;
 //! #[cfg_attr(feature = "async-trait", ractor::async_trait)]
-//! impl Actor for ExampleWorker {
-//!     type Msg = WorkerMessage<(), ExampleMessage>;
-//!     type State = WorkerStartContext<(), ExampleMessage, ()>;
-//!     type Arguments = WorkerStartContext<(), ExampleMessage, ()>;
+//! impl Worker for ExampleWorker {
+//!     type Key = ();
+//!     type Message = ExampleMessage;
+//!     type State = ();
+//!     type Arguments = ();
 //!     async fn pre_start(
 //!         &self,
-//!         _myself: ActorRef<Self::Msg>,
+//!         wid: WorkerId,
+//!         factory: &ActorRef<FactoryMessage<(), ExampleMessage>>,
 //!         startup_context: Self::Arguments,
 //!     ) -> Result<Self::State, ActorProcessingErr> {
 //!         Ok(startup_context)
 //!     }
 //!     async fn handle(
 //!         &self,
-//!         _myself: ActorRef<Self::Msg>,
-//!         message: Self::Msg,
-//!         state: &mut Self::State,
+//!         wid: WorkerId,
+//!         factory: &ActorRef<FactoryMessage<(), ExampleMessage>>,
+//!         Job {msg, key, ..}: Job<(), ExampleMessage>,
+//!         _state: &mut Self::State,
 //!     ) -> Result<(), ActorProcessingErr> {
-//!         match message {
-//!             WorkerMessage::FactoryPing(time) => {
-//!                 // This is a message which all factory workers **must**
-//!                 // adhere to. It is a background processing message from the
-//!                 // factory which is used for (a) metrics and (b) detecting
-//!                 // stuck workers, i.e. workers which aren't making progress
-//!                 // processing their messages
-//!                 state
-//!                     .factory
-//!                     .cast(FactoryMessage::WorkerPong(state.wid, time.elapsed()))?;
+//!         // Actual business logic that we want to parallelize
+//!         tracing::trace!("Worker {} received {:?}", wid, msg);
+//!         match msg {
+//!             ExampleMessage::PrintValue(value) => {
+//!                 tracing::info!("Worker {} printing value {value}", wid);
 //!             }
-//!             WorkerMessage::Dispatch(job) => {
-//!                 // Actual business logic that we want to parallelize
-//!                 tracing::trace!("Worker {} received {:?}", state.wid, job.msg);
-//!                 match job.msg {
-//!                     ExampleMessage::PrintValue(value) => {
-//!                         tracing::info!("Worker {} printing value {value}", state.wid);
-//!                     }
-//!                     ExampleMessage::EchoValue(value, reply) => {
-//!                         tracing::info!("Worker {} echoing value {value}", state.wid);
-//!                         let _ = reply.send(value);
-//!                     }
-//!                 }
-//!                 // job finished, on success or err we report back to the factory
-//!                 state
-//!                     .factory
-//!                     .cast(FactoryMessage::Finished(state.wid, job.key))?;
+//!             ExampleMessage::EchoValue(value, reply) => {
+//!                 tracing::info!("Worker {} echoing value {value}", wid);
+//!                 let _ = reply.send(value);
 //!             }
 //!         }
-//!         Ok(())
+//!         Ok(key)
 //!     }
 //! }
 //! /// Used by the factory to build new [ExampleWorker]s.
@@ -200,7 +185,7 @@ pub use factoryimpl::{Factory, FactoryArguments, FactoryArgumentsBuilder};
 pub use job::{Job, JobKey, JobOptions, MessageRetryStrategy, RetriableMessage};
 pub use lifecycle::FactoryLifecycleHooks;
 pub use worker::{
-    DeadMansSwitchConfiguration, WorkerBuilder, WorkerCapacityController, WorkerMessage,
+    DeadMansSwitchConfiguration, Worker, WorkerBuilder, WorkerCapacityController, WorkerMessage,
     WorkerProperties, WorkerStartContext,
 };
 
diff --git a/ractor/src/factory/queues.rs b/ractor/src/factory/queues.rs
index 11a6a57..4073248 100644
--- a/ractor/src/factory/queues.rs
+++ b/ractor/src/factory/queues.rs
@@ -423,10 +423,7 @@ mod tests {
             key: 99,
             accepted: None,
             msg: (),
-            options: JobOptions {
-                ttl: Some(Duration::from_millis(1)),
-                ..Default::default()
-            },
+            options: JobOptions::new(Some(Duration::from_millis(1))),
         });
 
         let oldest = queue.discard_oldest();
@@ -480,10 +477,7 @@ mod tests {
             key: 99,
             accepted: None,
             msg: (),
-            options: JobOptions {
-                ttl: Some(Duration::from_millis(1)),
-                ..Default::default()
-            },
+            options: JobOptions::new(Some(Duration::from_millis(1))),
         });
 
         // should discard lowest pri first
diff --git a/ractor/src/factory/worker.rs b/ractor/src/factory/worker.rs
index 92dc3a9..6d68ed8 100644
--- a/ractor/src/factory/worker.rs
+++ b/ractor/src/factory/worker.rs
@@ -10,10 +10,11 @@ use std::fmt::Debug;
 use std::sync::Arc;
 
 use bon::Builder;
+use tracing::Instrument;
 
 use crate::concurrency::{Duration, Instant, JoinHandle};
 use crate::{Actor, ActorId, ActorProcessingErr};
-use crate::{ActorRef, Message, MessagingErr};
+use crate::{ActorCell, ActorRef, Message, MessagingErr, SupervisionEvent};
 
 use super::discard::{DiscardMode, WorkerDiscardSettings};
 use super::stats::FactoryStatsLayer;
@@ -36,6 +37,459 @@ pub struct DeadMansSwitchConfiguration {
     pub kill_worker: bool,
 }
 
+/// A factory worker trait, which is a basic wrapper around
+/// actor logic, with predefined type information specific to workers
+///
+/// IMPORTANT: Workers are actors at their core principal, but with
+/// somewhat customized logic. This logic assists in tracking worker health,
+/// processing messages in a load-balanced manner, and managing necessary
+/// start automatically without copying the code repeatedly.
+///
+/// This trait implements as much of the custom wrapping logic as possible
+/// without breaking the factory <-> worker API requirement. If you so wish
+/// you can fully specify the actor properties instead of using this
+/// assistance trait.
+#[cfg_attr(feature = "async-trait", crate::async_trait)]
+pub trait Worker: Send + Sync + 'static {
+    /// The worker's job-key type
+    type Key: JobKey;
+    /// The worker's message type
+    type Message: Message;
+    /// The optional startup arguments for the worker (use `()` to ignore)
+    type Arguments: Message;
+    /// The worker's internal state
+    type State: crate::State;
+
+    /// Invoked when a worker is being started by the system.
+    ///
+    /// Any initialization inherent to the actor's role should be
+    /// performed here hence why it returns the initial state.
+    ///
+    /// Panics in `pre_start` do not invoke the
+    /// supervision strategy and the actor won't be started. The `spawn`
+    /// will return an error to the caller
+    ///
+    /// * `wid` - The id of this worker in the factory
+    /// * `factory` - The handle to the factory that owns and manages this worker
+    /// * `args` - Arguments that are passed in the spawning of the worker which are
+    /// necessary to construct the initial state
+    ///
+    /// Returns an initial [Worker::State] to bootstrap the actor
+    #[cfg(not(feature = "async-trait"))]
+    fn pre_start(
+        &self,
+        wid: WorkerId,
+        factory: &ActorRef<FactoryMessage<Self::Key, Self::Message>>,
+        args: Self::Arguments,
+    ) -> impl Future<Output = Result<Self::State, ActorProcessingErr>> + Send;
+
+    /// Invoked when a worker is being started by the system.
+    ///
+    /// Any initialization inherent to the actor's role should be
+    /// performed here hence why it returns the initial state.
+    ///
+    /// Panics in `pre_start` do not invoke the
+    /// supervision strategy and the actor won't be started. The `spawn`
+    /// will return an error to the caller
+    ///
+    /// * `wid` - The id of this worker in the factory
+    /// * `factory` - The handle to the factory that owns and manages this worker
+    /// * `args` - Arguments that are passed in the spawning of the worker which are
+    /// necessary to construct the initial state
+    ///
+    /// Returns an initial [Worker::State] to bootstrap the actor
+    #[cfg(feature = "async-trait")]
+    async fn pre_start(
+        &self,
+        wid: WorkerId,
+        factory: &ActorRef<FactoryMessage<Self::Key, Self::Message>>,
+        args: Self::Arguments,
+    ) -> Result<Self::State, ActorProcessingErr>;
+
+    /// Invoked after an actor has started.
+    ///
+    /// Any post initialization can be performed here, such as writing
+    /// to a log file, emitting metrics.
+    ///
+    /// * `wid` - The id of this worker in the factory
+    /// * `factory` - The handle to the factory that owns and manages this worker
+    /// * `state` - The worker's internal state, which is mutable and owned by the worker
+    ///
+    /// Panics in `post_start` follow the supervision strategy.
+    #[allow(unused_variables)]
+    #[cfg(not(feature = "async-trait"))]
+    fn post_start(
+        &self,
+        wid: WorkerId,
+        factory: &ActorRef<FactoryMessage<Self::Key, Self::Message>>,
+        state: &mut Self::State,
+    ) -> impl Future<Output = Result<(), ActorProcessingErr>> + Send {
+        async { Ok(()) }
+    }
+    /// Invoked after an actor has started.
+    ///
+    /// Any post initialization can be performed here, such as writing
+    /// to a log file, emitting metrics.
+    ///
+    /// * `wid` - The id of this worker in the factory
+    /// * `factory` - The handle to the factory that owns and manages this worker
+    /// * `state` - The worker's internal state, which is mutable and owned by the worker
+    ///
+    /// Panics in `post_start` follow the supervision strategy.
+    #[allow(unused_variables)]
+    #[cfg(feature = "async-trait")]
+    async fn post_start(
+        &self,
+        wid: WorkerId,
+        factory: &ActorRef<FactoryMessage<Self::Key, Self::Message>>,
+        state: &mut Self::State,
+    ) -> Result<(), ActorProcessingErr> {
+        Ok(())
+    }
+
+    /// Invoked after an actor has been stopped to perform final cleanup. In the
+    /// event the actor is terminated with killed or has self-panicked,
+    /// `post_stop` won't be called.
+    ///
+    /// * `wid` - The id of this worker in the factory
+    /// * `factory` - The handle to the factory that owns and manages this worker
+    /// * `state` - The worker's internal state, which is mutable and owned by the worker
+    ///
+    /// Panics in `post_stop` follow the supervision strategy.
+    #[allow(unused_variables)]
+    #[cfg(not(feature = "async-trait"))]
+    fn post_stop(
+        &self,
+        wid: WorkerId,
+        factory: &ActorRef<FactoryMessage<Self::Key, Self::Message>>,
+        state: &mut Self::State,
+    ) -> impl Future<Output = Result<(), ActorProcessingErr>> + Send {
+        async { Ok(()) }
+    }
+    /// Invoked after an actor has been stopped to perform final cleanup. In the
+    /// event the actor is terminated with killed or has self-panicked,
+    /// `post_stop` won't be called.
+    ///
+    /// * `wid` - The id of this worker in the factory
+    /// * `factory` - The handle to the factory that owns and manages this worker
+    /// * `state` - The worker's internal state, which is mutable and owned by the worker
+    ///
+    /// Panics in `post_stop` follow the supervision strategy.
+    #[allow(unused_variables)]
+    #[cfg(feature = "async-trait")]
+    async fn post_stop(
+        &self,
+        wid: WorkerId,
+        factory: &ActorRef<FactoryMessage<Self::Key, Self::Message>>,
+        state: &mut Self::State,
+    ) -> Result<(), ActorProcessingErr> {
+        Ok(())
+    }
+
+    /// Handle the incoming message from the event processing loop. Unhandled panickes will be
+    /// captured and sent to the supervisor(s)
+    ///
+    /// * `wid` - The id of this worker in the factory
+    /// * `factory` - The handle to the factory that owns and manages this worker
+    /// * `job` - The [Job] which this worker should process
+    /// * `state` - The worker's internal state, which is mutable and owned by the worker
+    ///
+    /// Returns the [Job::key] upon success or the error on failure
+    #[allow(unused_variables)]
+    #[cfg(not(feature = "async-trait"))]
+    fn handle(
+        &self,
+        wid: WorkerId,
+        factory: &ActorRef<FactoryMessage<Self::Key, Self::Message>>,
+        job: Job<Self::Key, Self::Message>,
+        state: &mut Self::State,
+    ) -> impl Future<Output = Result<Self::Key, ActorProcessingErr>> + Send {
+        async { Ok(job.key) }
+    }
+
+    /// Handle the incoming message from the event processing loop. Unhandled panickes will be
+    /// captured and sent to the supervisor(s)
+    ///
+    /// * `wid` - The id of this worker in the factory
+    /// * `factory` - The handle to the factory that owns and manages this worker
+    /// * `job` - The [Job] which this worker should process
+    /// * `state` - The worker's internal state, which is mutable and owned by the worker
+    ///
+    /// Returns the [Job::key] upon success or the error on failure
+    #[allow(unused_variables)]
+    #[cfg(feature = "async-trait")]
+    async fn handle(
+        &self,
+        wid: WorkerId,
+        factory: &ActorRef<FactoryMessage<Self::Key, Self::Message>>,
+        job: Job<Self::Key, Self::Message>,
+        state: &mut Self::State,
+    ) -> Result<Self::Key, ActorProcessingErr> {
+        Ok(job.key)
+    }
+
+    /// Handle the incoming supervision event. Unhandled panics will be captured and
+    /// sent the the supervisor(s). The default supervision behavior is to exit the
+    /// supervisor on any child exit. To override this behavior, implement this function.
+    ///
+    /// * `myself` - A reference to this actor's ActorCell
+    /// * `message` - The message to process
+    /// * `state` - A mutable reference to the internal actor's state
+    #[allow(unused_variables)]
+    #[cfg(not(feature = "async-trait"))]
+    fn handle_supervisor_evt(
+        &self,
+        myself: ActorCell,
+        message: SupervisionEvent,
+        state: &mut Self::State,
+    ) -> impl Future<Output = Result<(), ActorProcessingErr>> + Send {
+        async move {
+            match message {
+                SupervisionEvent::ActorTerminated(who, _, _)
+                | SupervisionEvent::ActorFailed(who, _) => {
+                    myself.stop(None);
+                }
+                _ => {}
+            }
+            Ok(())
+        }
+    }
+    /// Handle the incoming supervision event. Unhandled panics will be captured and
+    /// sent the the supervisor(s). The default supervision behavior is to exit the
+    /// supervisor on any child exit. To override this behavior, implement this function.
+    ///
+    /// * `myself` - A reference to this actor's ActorCell
+    /// * `message` - The message to process
+    /// * `state` - A mutable reference to the internal actor's state
+    #[allow(unused_variables)]
+    #[cfg(feature = "async-trait")]
+    async fn handle_supervisor_evt(
+        &self,
+        myself: ActorCell,
+        message: SupervisionEvent,
+        state: &mut Self::State,
+    ) -> Result<(), ActorProcessingErr> {
+        match message {
+            SupervisionEvent::ActorTerminated(who, _, _)
+            | SupervisionEvent::ActorFailed(who, _) => {
+                myself.stop(None);
+            }
+            _ => {}
+        }
+        Ok(())
+    }
+}
+
+/// The inner state of the wrapped [Worker] but held privately in trust
+/// for the [Worker] implementation
+#[doc(hidden)]
+pub struct WorkerState<TWorker: Worker> {
+    factory: ActorRef<FactoryMessage<TWorker::Key, TWorker::Message>>,
+    wid: WorkerId,
+    state: TWorker::State,
+}
+
+impl<TWorker: Worker> std::fmt::Debug for WorkerState<TWorker> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "WorkerState")
+    }
+}
+
+#[cfg_attr(feature = "async-trait", crate::async_trait)]
+impl<T> Actor for T
+where
+    T: Worker,
+{
+    type Msg = WorkerMessage<<Self as Worker>::Key, <Self as Worker>::Message>;
+    type Arguments = WorkerStartContext<
+        <Self as Worker>::Key,
+        <Self as Worker>::Message,
+        <Self as Worker>::Arguments,
+    >;
+    type State = WorkerState<Self>;
+
+    #[cfg(feature = "async-trait")]
+    async fn pre_start(
+        &self,
+        _: ActorRef<Self::Msg>,
+        WorkerStartContext {
+            wid,
+            factory,
+            custom_start,
+        }: Self::Arguments,
+    ) -> Result<Self::State, ActorProcessingErr> {
+        let inner_state = <Self as Worker>::pre_start(self, wid, &factory, custom_start).await?;
+        Ok(Self::State {
+            wid,
+            factory,
+            state: inner_state,
+        })
+    }
+
+    #[cfg(not(feature = "async-trait"))]
+    fn pre_start(
+        &self,
+        _: ActorRef<Self::Msg>,
+        WorkerStartContext {
+            wid,
+            factory,
+            custom_start,
+        }: Self::Arguments,
+    ) -> impl Future<Output = Result<Self::State, ActorProcessingErr>> + Send {
+        let inner_state = <Self as Worker>::pre_start(&self, wid, &factory, custom_start).await?;
+        async {
+            Ok(Self::State {
+                wid,
+                factory,
+                state: inner_state,
+            })
+        }
+    }
+
+    #[cfg(not(feature = "async-trait"))]
+    fn post_start(
+        &self,
+        _: ActorRef<Self::Msg>,
+        state: &mut Self::State,
+    ) -> impl Future<Output = Result<(), ActorProcessingErr>> + Send {
+        async {
+            <Self as Worker>::post_start(&self, state.wid, &state.factory, &mut state.state).await
+        }
+    }
+
+    #[cfg(feature = "async-trait")]
+    async fn post_start(
+        &self,
+        _: ActorRef<Self::Msg>,
+        state: &mut Self::State,
+    ) -> Result<(), ActorProcessingErr> {
+        <Self as Worker>::post_start(self, state.wid, &state.factory, &mut state.state).await
+    }
+
+    #[cfg(not(feature = "async-trait"))]
+    fn post_stop(
+        &self,
+        _: ActorRef<Self::Msg>,
+        state: &mut Self::State,
+    ) -> impl Future<Output = Result<(), ActorProcessingErr>> + Send {
+        async {
+            <Self as Worker>::post_stop(&self, state.wid, &state.factory, &mut state.state).await
+        }
+    }
+
+    #[cfg(feature = "async-trait")]
+    async fn post_stop(
+        &self,
+        _: ActorRef<Self::Msg>,
+        state: &mut Self::State,
+    ) -> Result<(), ActorProcessingErr> {
+        <Self as Worker>::post_stop(self, state.wid, &state.factory, &mut state.state).await
+    }
+
+    #[cfg(not(feature = "async-trait"))]
+    fn handle(
+        &self,
+        _: ActorRef<Self::Msg>,
+        message: Self::Msg,
+        state: &mut Self::State,
+    ) -> impl Future<Output = Result<(), ActorProcessingErr>> + Send {
+        async {
+            match message {
+                WorkerMessage::FactoryPing(time) => {
+                    tracing::trace!("Worker {} - ping", state.wid);
+
+                    state
+                        .factory
+                        .cast(FactoryMessage::WorkerPong(state.wid, time.elapsed()))?;
+                }
+                WorkerMessage::Dispatch(mut job) => {
+                    let key = if let Some(span) = job.options.span.take() {
+                        <Self as Worker>::handle(
+                            &self,
+                            state.wid,
+                            &state.factory,
+                            job,
+                            &mut state.state,
+                        )
+                        .instrument(span)
+                        .await
+                    } else {
+                        <Self as Worker>::handle(
+                            &self,
+                            state.wid,
+                            &state.factory,
+                            job,
+                            &mut state.state,
+                        )
+                        .await
+                    }?;
+                    state
+                        .factory
+                        .cast(FactoryMessage::Finished(state.wid, key))?;
+                    Ok(())
+                }
+            }
+        }
+    }
+
+    #[cfg(feature = "async-trait")]
+    async fn handle(
+        &self,
+        _: ActorRef<Self::Msg>,
+        message: Self::Msg,
+        state: &mut Self::State,
+    ) -> Result<(), ActorProcessingErr> {
+        match message {
+            WorkerMessage::FactoryPing(time) => {
+                tracing::trace!("Worker {} - ping", state.wid);
+
+                state
+                    .factory
+                    .cast(FactoryMessage::WorkerPong(state.wid, time.elapsed()))?;
+                Ok(())
+            }
+            WorkerMessage::Dispatch(mut job) => {
+                let key = if let Some(span) = job.options.take_span() {
+                    <Self as Worker>::handle(self, state.wid, &state.factory, job, &mut state.state)
+                        .instrument(span)
+                        .await
+                } else {
+                    <Self as Worker>::handle(self, state.wid, &state.factory, job, &mut state.state)
+                        .await
+                }?;
+                state
+                    .factory
+                    .cast(FactoryMessage::Finished(state.wid, key))?;
+                Ok(())
+            }
+        }
+    }
+
+    #[cfg(not(feature = "async-trait"))]
+    fn handle_supervisor_evt(
+        &self,
+        myself: ActorRef<Self::Msg>,
+        message: SupervisionEvent,
+        state: &mut Self::State,
+    ) -> impl Future<Output = Result<(), ActorProcessingErr>> + Send {
+        async move {
+            <Self as Worker>::handle_supervisor_evt(&self, myself.into(), message, &mut state.state)
+                .await
+        }
+    }
+
+    #[cfg(feature = "async-trait")]
+    async fn handle_supervisor_evt(
+        &self,
+        myself: ActorRef<Self::Msg>,
+        message: SupervisionEvent,
+        state: &mut Self::State,
+    ) -> Result<(), ActorProcessingErr> {
+        <Self as Worker>::handle_supervisor_evt(self, myself.into(), message, &mut state.state)
+            .await
+    }
+}
+
 /// The [super::Factory] is responsible for spawning workers
 /// and re-spawning workers under failure scenarios. This means that
 /// it needs to understand how to build workers. The [WorkerBuilder]
diff --git a/ractor_cluster/Cargo.toml b/ractor_cluster/Cargo.toml
index 716cc1d..51938f2 100644
--- a/ractor_cluster/Cargo.toml
+++ b/ractor_cluster/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "ractor_cluster"
-version = "0.13.5"
-authors = ["Sean Lawlor", "Evan Au", "Dillon George"]
+version = "0.14.0"
+authors = ["Sean Lawlor <slawlor>"]
 description = "Distributed cluster environment of Ractor actors"
 documentation = "https://docs.rs/ractor"
 license = "MIT"
@@ -24,8 +24,8 @@ prost-build = { version = "0.13" }
 bytes = { version = "1" }
 prost = { version = "0.13" }
 prost-types = { version = "0.13" }
-ractor = { version = "0.13.0", features = ["cluster"], path = "../ractor" }
-ractor_cluster_derive = { version = "0.13.0", path = "../ractor_cluster_derive" }
+ractor = { version = "0.14.0", features = ["cluster"], path = "../ractor" }
+ractor_cluster_derive = { version = "0.14.0", path = "../ractor_cluster_derive" }
 rand = "0.8"
 sha2 = "0.10"
 tokio = { version = "1.30", features = ["rt", "time", "sync", "macros", "net", "io-util", "tracing"]}
diff --git a/ractor_cluster_derive/Cargo.toml b/ractor_cluster_derive/Cargo.toml
index 5d18794..cf0cccb 100644
--- a/ractor_cluster_derive/Cargo.toml
+++ b/ractor_cluster_derive/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "ractor_cluster_derive"
-version = "0.13.5"
-authors = ["Sean Lawlor <seanlawlor@fb.com>"]
+version = "0.14.0"
+authors = ["Sean Lawlor <slawlor>"]
 description = "Derives for ractor_cluster"
 license = "MIT"
 edition = "2021"


Test patch:
diff --git a/ractor/src/factory/tests/basic.rs b/ractor/src/factory/tests/basic.rs
index bdbd1e1..07c5e53 100644
--- a/ractor/src/factory/tests/basic.rs
+++ b/ractor/src/factory/tests/basic.rs
@@ -55,14 +55,16 @@ struct TestWorker {
 }
 
 #[cfg_attr(feature = "async-trait", crate::async_trait)]
-impl Actor for TestWorker {
-    type Msg = WorkerMessage<TestKey, TestMessage>;
-    type State = Self::Arguments;
-    type Arguments = WorkerStartContext<TestKey, TestMessage, ()>;
+impl Worker for TestWorker {
+    type Key = TestKey;
+    type Message = TestMessage;
+    type State = ();
+    type Arguments = ();
 
     async fn pre_start(
         &self,
-        _myself: ActorRef<Self::Msg>,
+        _wid: WorkerId,
+        _factory: &ActorRef<FactoryMessage<TestKey, TestMessage>>,
         startup_context: Self::Arguments,
     ) -> Result<Self::State, ActorProcessingErr> {
         Ok(startup_context)
@@ -70,32 +72,20 @@ impl Actor for TestWorker {
 
     async fn handle(
         &self,
-        _myself: ActorRef<Self::Msg>,
-        message: Self::Msg,
-        state: &mut Self::State,
-    ) -> Result<(), ActorProcessingErr> {
-        match message {
-            WorkerMessage::FactoryPing(time) => {
-                state
-                    .factory
-                    .cast(FactoryMessage::WorkerPong(state.wid, time.elapsed()))?;
-            }
-            WorkerMessage::Dispatch(job) => {
-                tracing::debug!("Worker received {:?}", job.msg);
+        _wid: WorkerId,
+        _factory: &ActorRef<FactoryMessage<TestKey, TestMessage>>,
+        Job { msg, key, .. }: Job<Self::Key, Self::Message>,
+        _state: &mut Self::State,
+    ) -> Result<TestKey, ActorProcessingErr> {
+        tracing::debug!("Worker received {:?}", msg);
 
-                self.counter.fetch_add(1, Ordering::Relaxed);
-
-                if let Some(timeout_ms) = self.slow {
-                    crate::concurrency::sleep(Duration::from_millis(timeout_ms)).await;
-                }
+        self.counter.fetch_add(1, Ordering::Relaxed);
 
-                // job finished, on success or err we report back to the factory
-                state
-                    .factory
-                    .cast(FactoryMessage::Finished(state.wid, job.key))?;
-            }
+        if let Some(timeout_ms) = self.slow {
+            crate::concurrency::sleep(Duration::from_millis(timeout_ms)).await;
         }
-        Ok(())
+
+        Ok(key)
     }
 }
 
diff --git a/ractor/src/factory/tests/dynamic_pool.rs b/ractor/src/factory/tests/dynamic_pool.rs
index d8cd68c..688ad5d 100644
--- a/ractor/src/factory/tests/dynamic_pool.rs
+++ b/ractor/src/factory/tests/dynamic_pool.rs
@@ -48,14 +48,16 @@ struct TestWorker {
 impl crate::Message for TestMessage {}
 
 #[cfg_attr(feature = "async-trait", crate::async_trait)]
-impl Actor for TestWorker {
-    type Msg = WorkerMessage<TestKey, TestMessage>;
+impl Worker for TestWorker {
+    type Key = TestKey;
+    type Message = TestMessage;
     type State = Self::Arguments;
-    type Arguments = WorkerStartContext<TestKey, TestMessage, ()>;
+    type Arguments = ();
 
     async fn pre_start(
         &self,
-        _myself: ActorRef<Self::Msg>,
+        _wid: WorkerId,
+        _factory: &ActorRef<FactoryMessage<Self::Key, Self::Message>>,
         startup_context: Self::Arguments,
     ) -> Result<Self::State, ActorProcessingErr> {
         Ok(startup_context)
@@ -63,28 +65,15 @@ impl Actor for TestWorker {
 
     async fn handle(
         &self,
-        _myself: ActorRef<Self::Msg>,
-        message: Self::Msg,
-        state: &mut Self::State,
-    ) -> Result<(), ActorProcessingErr> {
-        match message {
-            WorkerMessage::FactoryPing(time) => {
-                state
-                    .factory
-                    .cast(FactoryMessage::WorkerPong(state.wid, time.elapsed()))?;
-            }
-            WorkerMessage::Dispatch(job) => {
-                tracing::debug!("Worker received {:?}", job.msg);
-
-                self.id_map.insert(state.wid);
-
-                // job finished, on success or err we report back to the factory
-                state
-                    .factory
-                    .cast(FactoryMessage::Finished(state.wid, job.key))?;
-            }
-        }
-        Ok(())
+        wid: WorkerId,
+        _factory: &ActorRef<FactoryMessage<Self::Key, Self::Message>>,
+        Job { msg, key, .. }: Job<Self::Key, Self::Message>,
+        _state: &mut Self::State,
+    ) -> Result<Self::Key, ActorProcessingErr> {
+        tracing::debug!("Worker received {:?}", msg);
+
+        self.id_map.insert(wid);
+        Ok(key)
     }
 }