Made durable generic over a context type (#20)

* wip

* made durable generic over a context type

* renamed app context to state

* made TaskContext generic over app state

Author: virajmehta

benches/common/tasks.rs

@@ -9,12 +9,12 @@ use serde::{Deserialize, Serialize};
 pub struct NoOpTask;
 
 #[async_trait]
-impl Task for NoOpTask {
+impl Task<()> for NoOpTask {
     const NAME: &'static str = "bench-noop";
     type Params = ();
     type Output = ();
 
-    async fn run(_params: Self::Params, _ctx: TaskContext) -> TaskResult<Self::Output> {
+    async fn run(_params: Self::Params, _ctx: TaskContext, _state: ()) -> TaskResult<Self::Output> {
         Ok(())
     }
 }
@@ -33,12 +33,12 @@ pub struct QuickParams {
 }
 
 #[async_trait]
-impl Task for QuickTask {
+impl Task<()> for QuickTask {
     const NAME: &'static str = "bench-quick";
     type Params = QuickParams;
     type Output = u32;
 
-    async fn run(params: Self::Params, _ctx: TaskContext) -> TaskResult<Self::Output> {
+    async fn run(params: Self::Params, _ctx: TaskContext, _state: ()) -> TaskResult<Self::Output> {
         Ok(params.task_num)
     }
 }
@@ -57,12 +57,16 @@ pub struct MultiStepParams {
 }
 
 #[async_trait]
-impl Task for MultiStepBenchTask {
+impl Task<()> for MultiStepBenchTask {
     const NAME: &'static str = "bench-multi-step";
     type Params = MultiStepParams;
     type Output = u32;
 
-    async fn run(params: Self::Params, mut ctx: TaskContext) -> TaskResult<Self::Output> {
+    async fn run(
+        params: Self::Params,
+        mut ctx: TaskContext,
+        _state: (),
+    ) -> TaskResult<Self::Output> {
         for i in 0..params.num_steps {
             let _: u32 = ctx
                 .step(&format!("step-{}", i), || async move { Ok(i) })
@@ -86,12 +90,16 @@ pub struct LargePayloadParams {
 }
 
 #[async_trait]
-impl Task for LargePayloadBenchTask {
+impl Task<()> for LargePayloadBenchTask {
     const NAME: &'static str = "bench-large-payload";
     type Params = LargePayloadParams;
     type Output = usize;
 
-    async fn run(params: Self::Params, mut ctx: TaskContext) -> TaskResult<Self::Output> {
+    async fn run(
+        params: Self::Params,
+        mut ctx: TaskContext,
+        _state: (),
+    ) -> TaskResult<Self::Output> {
         let payload = "x".repeat(params.payload_size);
         let _: String = ctx
             .step("large-step", || async move { Ok(payload) })

src/client.rs

@@ -57,37 +57,67 @@ use crate::worker::Worker;
 /// - Emit events with [`emit_event`](Self::emit_event)
 /// - Cancel tasks with [`cancel_task`](Self::cancel_task)
 ///
+/// # Type Parameter
+///
+/// * `State` - Application state type passed to task handlers. Use `()` if you
+///   don't need any state. The state must implement `Clone + Send + Sync + 'static`.
+///
 /// # Example
 ///
 /// ```ignore
+/// // Without state (default)
 /// let client = Durable::builder()
 ///     .database_url("postgres://localhost/myapp")
 ///     .queue_name("tasks")
 ///     .build()
 ///     .await?;
 ///
+/// // With application state
+/// #[derive(Clone)]
+/// struct AppState {
+///     http_client: reqwest::Client,
+/// }
+///
+/// let app_state = AppState { http_client: reqwest::Client::new() };
+/// let client = Durable::builder()
+///     .database_url("postgres://localhost/myapp")
+///     .queue_name("tasks")
+///     .build_with_state(app_state)
+///     .await?;
+///
 /// client.register::<MyTask>().await;
 /// client.spawn::<MyTask>(params).await?;
 /// ```
-pub struct Durable {
+pub struct Durable<State = ()>
+where
+    State: Clone + Send + Sync + 'static,
+{
     pool: PgPool,
     owns_pool: bool,
     queue_name: String,
     default_max_attempts: u32,
-    registry: Arc<RwLock<TaskRegistry>>,
+    registry: Arc<RwLock<TaskRegistry<State>>>,
+    state: State,
 }
 
 /// Builder for configuring a [`Durable`] client.
 ///
 /// # Example
 ///
 /// ```ignore
+/// // Without state
 /// let client = Durable::builder()
 ///     .database_url("postgres://localhost/myapp")
 ///     .queue_name("orders")
 ///     .default_max_attempts(3)
 ///     .build()
 ///     .await?;
+///
+/// // With state
+/// let client = Durable::builder()
+///     .database_url("postgres://localhost/myapp")
+///     .build_with_state(my_app_state)
+///     .await?;
 /// ```
 pub struct DurableBuilder {
     database_url: Option<String>,
@@ -130,8 +160,43 @@ impl DurableBuilder {
         self
     }
 
-    /// Build the Durable client
-    pub async fn build(self) -> anyhow::Result<Durable> {
+    /// Build the Durable client without application state.
+    ///
+    /// Use this when your tasks don't need access to shared resources
+    /// like HTTP clients or database pools.
+    pub async fn build(self) -> anyhow::Result<Durable<()>> {
+        self.build_with_state(()).await
+    }
+
+    /// Build the Durable client with application state.
+    ///
+    /// The state will be cloned and passed to each task execution.
+    /// Use this to provide shared resources like HTTP clients, database pools,
+    /// or other application state to your tasks.
+    ///
+    /// # Example
+    ///
+    /// ```ignore
+    /// #[derive(Clone)]
+    /// struct AppState {
+    ///     http_client: reqwest::Client,
+    ///     db_pool: PgPool,
+    /// }
+    ///
+    /// let state = AppState {
+    ///     http_client: reqwest::Client::new(),
+    ///     db_pool: pool.clone(),
+    /// };
+    ///
+    /// let client = Durable::builder()
+    ///     .database_url("postgres://localhost/myapp")
+    ///     .build_with_state(state)
+    ///     .await?;
+    /// ```
+    pub async fn build_with_state<State>(self, state: State) -> anyhow::Result<Durable<State>>
+    where
+        State: Clone + Send + Sync + 'static,
+    {
         let (pool, owns_pool) = if let Some(pool) = self.pool {
             (pool, false)
         } else {
@@ -148,6 +213,7 @@ impl DurableBuilder {
             queue_name: self.queue_name,
             default_max_attempts: self.default_max_attempts,
             registry: Arc::new(RwLock::new(HashMap::new())),
+            state,
         })
     }
 }
@@ -158,8 +224,8 @@ impl Default for DurableBuilder {
     }
 }
 
-impl Durable {
-    /// Create a new client with default settings
+impl Durable<()> {
+    /// Create a new client with default settings (no application state).
     pub async fn new(database_url: &str) -> anyhow::Result<Self> {
         DurableBuilder::new()
             .database_url(database_url)
@@ -171,7 +237,12 @@ impl Durable {
     pub fn builder() -> DurableBuilder {
         DurableBuilder::new()
     }
+}
 
+impl<State> Durable<State>
+where
+    State: Clone + Send + Sync + 'static,
+{
     /// Get a reference to the underlying connection pool
     pub fn pool(&self) -> &PgPool {
         &self.pool
@@ -182,21 +253,29 @@ impl Durable {
         &self.queue_name
     }
 
+    /// Get a reference to the application state
+    pub fn state(&self) -> &State {
+        &self.state
+    }
+
     /// Register a task type. Required before spawning or processing.
-    pub async fn register<T: Task>(&self) -> &Self {
+    pub async fn register<T: Task<State>>(&self) -> &Self {
         let mut registry = self.registry.write().await;
-        registry.insert(T::NAME.to_string(), Arc::new(TaskWrapper::<T>::new()));
+        registry.insert(
+            T::NAME.to_string(),
+            Arc::new(TaskWrapper::<T, State>::new()),
+        );
         self
     }
 
     /// Spawn a task (type-safe version)
-    pub async fn spawn<T: Task>(&self, params: T::Params) -> anyhow::Result<SpawnResult> {
+    pub async fn spawn<T: Task<State>>(&self, params: T::Params) -> anyhow::Result<SpawnResult> {
         self.spawn_with_options::<T>(params, SpawnOptions::default())
             .await
     }
 
     /// Spawn a task with options (type-safe version)
-    pub async fn spawn_with_options<T: Task>(
+    pub async fn spawn_with_options<T: Task<State>>(
         &self,
         params: T::Params,
         options: SpawnOptions,
@@ -240,7 +319,7 @@ impl Durable {
         params: T::Params,
     ) -> anyhow::Result<SpawnResult>
     where
-        T: Task,
+        T: Task<State>,
         E: Executor<'e, Database = Postgres>,
     {
         self.spawn_with_options_with::<T, E>(executor, params, SpawnOptions::default())
@@ -255,7 +334,7 @@ impl Durable {
         options: SpawnOptions,
     ) -> anyhow::Result<SpawnResult>
     where
-        T: Task,
+        T: Task<State>,
         E: Executor<'e, Database = Postgres>,
     {
         self.spawn_by_name_with(executor, T::NAME, serde_json::to_value(&params)?, options)
@@ -377,6 +456,7 @@ impl Durable {
             self.queue_name.clone(),
             self.registry.clone(),
             options,
+            self.state.clone(),
         )
         .await
     }

src/context.rs

@@ -3,6 +3,7 @@ use serde::{Serialize, de::DeserializeOwned};
 use serde_json::Value as JsonValue;
 use sqlx::PgPool;
 use std::collections::HashMap;
+use std::marker::PhantomData;
 use std::time::Duration;
 use uuid::Uuid;
 
@@ -27,12 +28,20 @@ use crate::worker::LeaseExtender;
 /// - **Lease management** via [`heartbeat`](Self::heartbeat) - Extend the task lease
 ///   for long-running operations
 ///
+/// # Type Parameter
+///
+/// * `State` - The application state type. This allows [`spawn`](Self::spawn) to
+///   automatically infer the correct state type for child tasks.
+///
 /// # Public Fields
 ///
 /// - `task_id` - Unique identifier for this task (use as idempotency key)
 /// - `run_id` - Identifier for the current execution attempt
 /// - `attempt` - Current attempt number (starts at 1)
-pub struct TaskContext {
+pub struct TaskContext<State = ()>
+where
+    State: Clone + Send + Sync + 'static,
+{
     /// Unique identifier for this task. Use this as an idempotency key for
     /// external API calls to achieve "exactly-once" semantics.
     pub task_id: Uuid,
@@ -57,6 +66,9 @@ pub struct TaskContext {
 
     /// Notifies the worker when the lease is extended via step() or heartbeat().
     lease_extender: LeaseExtender,
+
+    /// Phantom data to carry the State type parameter.
+    _state: PhantomData<State>,
 }
 
 /// Validate that a user-provided step name doesn't use reserved prefix.
@@ -69,7 +81,10 @@ fn validate_user_name(name: &str) -> TaskResult<()> {
     Ok(())
 }
 
-impl TaskContext {
+impl<State> TaskContext<State>
+where
+    State: Clone + Send + Sync + 'static,
+{
     /// Create a new TaskContext. Called by the worker before executing a task.
     /// Loads all existing checkpoints into the cache.
     pub(crate) async fn create(
@@ -105,6 +120,7 @@ impl TaskContext {
             checkpoint_cache: cache,
             step_counters: HashMap::new(),
             lease_extender,
+            _state: PhantomData,
         })
     }
 
@@ -431,12 +447,15 @@ impl TaskContext {
     /// let r1: ItemResult = ctx.join("item-1", h1).await?;
     /// let r2: ItemResult = ctx.join("item-2", h2).await?;
     /// ```
-    pub async fn spawn<T: Task>(
+    pub async fn spawn<T>(
         &mut self,
         name: &str,
         params: T::Params,
         options: crate::SpawnOptions,
-    ) -> TaskResult<TaskHandle<T::Output>> {
+    ) -> TaskResult<TaskHandle<T::Output>>
+    where
+        T: Task<State>,
+    {
         validate_user_name(name)?;
         let checkpoint_name = self.get_checkpoint_name(&format!("$spawn:{name}"));