diff --git a/docs/migrating-to-0.5.md b/docs/migrating-to-0.5.md index c8cc121..79bd5ab 100644 --- a/docs/migrating-to-0.5.md +++ b/docs/migrating-to-0.5.md @@ -165,7 +165,7 @@ The `finalize` and `on_cancel` hooks follow the same pattern: ```rust impl TypedExecutor for ThumbnailExec { - async fn finalize(&self, thumb: Thumbnail, ctx: &TaskContext) -> Result<(), TaskError> { + async fn finalize(&self, thumb: Thumbnail, _memo: (), ctx: &TaskContext) -> Result<(), TaskError> { // called after all children settle Ok(()) } diff --git a/docs/quick-start.md b/docs/quick-start.md index 699b1b2..566f022 100644 --- a/docs/quick-start.md +++ b/docs/quick-start.md @@ -277,7 +277,7 @@ impl TypedExecutor for MultipartUploader { } async fn finalize( - &self, upload: MultipartUpload, ctx: &TaskContext, + &self, upload: MultipartUpload, _memo: (), ctx: &TaskContext, ) -> Result<(), TaskError> { // Called after all children complete complete_multipart_upload(&upload).await diff --git a/migrations/009_memo.sql b/migrations/009_memo.sql new file mode 100644 index 0000000..8db0b72 --- /dev/null +++ b/migrations/009_memo.sql @@ -0,0 +1,3 @@ +-- Execute-to-finalize memo: typed state persisted between phases. +ALTER TABLE tasks ADD COLUMN memo BLOB DEFAULT NULL; +ALTER TABLE task_history ADD COLUMN memo BLOB DEFAULT NULL; diff --git a/src/domain.rs b/src/domain.rs index 6067f7b..4d33c7e 100644 --- a/src/domain.rs +++ b/src/domain.rs @@ -20,6 +20,8 @@ use std::pin::Pin; use std::sync::Arc; use std::time::Duration; +use serde::{de::DeserializeOwned, Serialize}; + use crate::module::{ExecutorOptions, ModuleExecutor, ModuleHandle}; use crate::priority::Priority; use crate::registry::{DomainTaskContext, ErasedExecutor, TaskContext, TaskExecutor}; @@ -181,19 +183,29 @@ pub struct TaskTypeOptions { /// } /// } /// ``` -pub trait TypedExecutor: Send + Sync + 'static { +pub trait TypedExecutor< + T: TypedTask, + Memo: Serialize + DeserializeOwned + Send + Sync + 'static = (), +>: Send + Sync + 'static +{ /// Primary execution. Called once per dispatch. + /// + /// Returns a `Memo` that will be persisted and passed to [`finalize()`](Self::finalize) + /// after all children complete. For the default `Memo = ()`, the return type + /// is `Result<(), TaskError>` — identical to the pre-memo API. fn execute<'a>( &'a self, payload: T, ctx: DomainTaskContext<'a, T::Domain>, - ) -> impl Future> + Send + 'a; + ) -> impl Future> + Send + 'a; /// Called when all child tasks spawned by this task have settled. + /// Receives the `Memo` returned by [`execute()`](Self::execute). /// Default: no-op. fn finalize<'a>( &'a self, _payload: T, + _memo: Memo, _ctx: DomainTaskContext<'a, T::Domain>, ) -> impl Future> + Send + 'a { async { Ok(()) } @@ -212,26 +224,46 @@ pub trait TypedExecutor: Send + Sync + 'static { // ── TypedExecutorAdapter ───────────────────────────────────────────── -/// Internal adapter that wraps a [`TypedExecutor`] into a [`TaskExecutor`] +/// Internal adapter that wraps a [`TypedExecutor`] into a [`TaskExecutor`] /// for the scheduler engine. /// -/// Handles payload deserialization before delegating to the typed executor. -struct TypedExecutorAdapter { +/// Handles payload deserialization and memo serialization/deserialization. +struct TypedExecutorAdapter { executor: E, - _marker: PhantomData T>, + _marker: PhantomData (T, M)>, } -impl> TaskExecutor for TypedExecutorAdapter { - async fn execute<'a>(&'a self, ctx: &'a TaskContext) -> Result<(), TaskError> { +impl TaskExecutor for TypedExecutorAdapter +where + T: TypedTask, + M: Serialize + DeserializeOwned + Send + Sync + 'static, + E: TypedExecutor, +{ + async fn execute<'a>(&'a self, ctx: &'a TaskContext) -> Result>, TaskError> { let payload: T = ctx.payload()?; let dctx = DomainTaskContext::::new(ctx); - self.executor.execute(payload, dctx).await + let memo = self.executor.execute(payload, dctx).await?; + + // Don't persist () — serialize to None. + if std::any::TypeId::of::() == std::any::TypeId::of::<()>() { + return Ok(None); + } + + let bytes = serde_json::to_vec(&memo) + .map_err(|e| TaskError::permanent(format!("memo serialization: {e}")))?; + Ok(Some(bytes)) } async fn finalize<'a>(&'a self, ctx: &'a TaskContext) -> Result<(), TaskError> { let payload: T = ctx.payload()?; + let memo: M = match &ctx.record().memo { + Some(bytes) => serde_json::from_slice(bytes) + .map_err(|e| TaskError::permanent(format!("memo deserialization: {e}")))?, + None => serde_json::from_value(serde_json::Value::Null) + .map_err(|e| TaskError::permanent(format!("memo deserialization: {e}")))?, + }; let dctx = DomainTaskContext::::new(ctx); - self.executor.finalize(payload, dctx).await + self.executor.finalize(payload, memo, dctx).await } async fn on_cancel<'a>(&'a self, ctx: &'a TaskContext) -> Result<(), TaskError> { @@ -241,6 +273,19 @@ impl> TaskExecutor for TypedExecutorAdapter(executor: E) -> Arc +where + T: TypedTask, + M: Serialize + DeserializeOwned + Send + Sync + 'static, + E: TypedExecutor, +{ + Arc::new(TypedExecutorAdapter { + executor, + _marker: PhantomData:: (T, M)>, + }) +} + // ── Domain ──────────────────────────────────────────────────────── /// A typed module builder that enforces the link between a [`DomainKey`], @@ -317,7 +362,36 @@ impl Domain { T: TypedTask, { let config = T::config(); - self.task_inner::(executor, config.ttl, config.retry_policy) + self.task_inner::( + erase_executor::(executor), + config.ttl, + config.retry_policy, + ) + } + + /// Register a typed executor that produces a memo in `execute()` which + /// is persisted and passed to `finalize()`. + /// + /// Both `T` and `Memo` are inferred from the executor's + /// `TypedExecutor` impl — turbofish is only needed when the + /// executor is generic over task types. + /// + /// # Example + /// + /// ```ignore + /// domain.task_memo(ScanL1Executor) + /// ``` + pub fn task_memo(self, executor: impl TypedExecutor) -> Self + where + T: TypedTask, + Memo: Serialize + DeserializeOwned + Send + Sync + 'static, + { + let config = T::config(); + self.task_inner::( + erase_executor::(executor), + config.ttl, + config.retry_policy, + ) } /// Register a typed executor with per-type option overrides. @@ -332,25 +406,35 @@ impl Domain { let config = T::config(); let ttl = options.ttl.or(config.ttl); let retry_policy = options.retry_policy.or(config.retry_policy); - self.task_inner::(executor, ttl, retry_policy) + self.task_inner::(erase_executor::(executor), ttl, retry_policy) } - fn task_inner( - mut self, - executor: impl TypedExecutor, - ttl: Option, - retry_policy: Option, + /// Like [`task_with()`](Self::task_with), but for executors that produce + /// a memo (see [`task_memo()`](Self::task_memo)). + pub fn task_with_memo( + self, + executor: impl TypedExecutor, + options: TaskTypeOptions, ) -> Self where - T: TypedTask, + T: TypedTask, + Memo: Serialize + DeserializeOwned + Send + Sync + 'static, { - let adapter = TypedExecutorAdapter { - executor, - _marker: PhantomData:: T>, - }; + let config = T::config(); + let ttl = options.ttl.or(config.ttl); + let retry_policy = options.retry_policy.or(config.retry_policy); + self.task_inner::(erase_executor::(executor), ttl, retry_policy) + } + + fn task_inner( + mut self, + executor: Arc, + ttl: Option, + retry_policy: Option, + ) -> Self { self.executors.push(ModuleExecutor { task_type: T::TASK_TYPE.to_string(), - executor: Arc::new(adapter) as Arc, + executor, options: ExecutorOptions { ttl, retry_policy }, }); self diff --git a/src/lib.rs b/src/lib.rs index adcc01b..667ac31 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -534,7 +534,7 @@ //! Ok(()) //! } //! -//! async fn finalize(&self, upload: MultipartUpload, ctx: DomainTaskContext<'_, Uploads>) -> Result<(), TaskError> { +//! async fn finalize(&self, upload: MultipartUpload, _memo: (), ctx: DomainTaskContext<'_, Uploads>) -> Result<(), TaskError> { //! // All parts uploaded — complete the multipart upload. //! complete_multipart(&upload).await?; //! Ok(()) diff --git a/src/registry/mod.rs b/src/registry/mod.rs index b25a996..02c3fd6 100644 --- a/src/registry/mod.rs +++ b/src/registry/mod.rs @@ -61,13 +61,14 @@ pub(crate) trait TaskExecutor: Send + Sync + 'static { /// - `ctx`: Execution context with the task record, cancellation token, /// and progress reporter. /// - /// On success, return `Ok(())`. Use [`TaskContext::record_read_bytes`] + /// On success, return `Ok(None)` or `Ok(Some(bytes))` with serialized + /// memo data to pass to `finalize()`. Use [`TaskContext::record_read_bytes`] /// and [`TaskContext::record_write_bytes`] to report IO during execution. /// On failure, return a [`TaskError`] indicating whether retry is appropriate. fn execute<'a>( &'a self, ctx: &'a TaskContext, - ) -> impl Future> + Send + 'a; + ) -> impl Future>, TaskError>> + Send + 'a; /// Called after all children of a parent task have completed. /// @@ -110,6 +111,9 @@ pub struct TaskTypeRegistry { type_retry_policies: HashMap, } +/// Serialized memo bytes returned by `execute_erased`. +type MemoBytes = Option>; + /// Object-safe wrapper around [`TaskExecutor`] for dynamic dispatch in the registry. /// /// This trait exists because RPITIT (`impl Future`) in `TaskExecutor` is not @@ -119,7 +123,7 @@ pub(crate) trait ErasedExecutor: Send + Sync + 'static { fn execute_erased<'a>( &'a self, ctx: &'a TaskContext, - ) -> std::pin::Pin> + Send + 'a>>; + ) -> std::pin::Pin> + Send + 'a>>; fn finalize_erased<'a>( &'a self, @@ -136,7 +140,7 @@ impl ErasedExecutor for T { fn execute_erased<'a>( &'a self, ctx: &'a TaskContext, - ) -> std::pin::Pin> + Send + 'a>> { + ) -> std::pin::Pin> + Send + 'a>> { Box::pin(self.execute(ctx)) } diff --git a/src/scheduler/spawn.rs b/src/scheduler/spawn.rs index 16845d8..895108e 100644 --- a/src/scheduler/spawn.rs +++ b/src/scheduler/spawn.rs @@ -111,7 +111,9 @@ pub(crate) async fn spawn_task( let result = match phase { ExecutionPhase::Execute => executor.execute_erased(&prepared.ctx).await, - ExecutionPhase::Finalize => executor.finalize_erased(&prepared.ctx).await, + ExecutionPhase::Finalize => { + executor.finalize_erased(&prepared.ctx).await.map(|()| None) + } // finalize doesn't produce a memo }; // Read IO bytes from the context tracker. @@ -121,11 +123,12 @@ pub(crate) async fn spawn_task( drop(prepared.ctx); match result { - Ok(()) => { + Ok(memo) => { completion::handle_success( &task, phase, &metrics, + memo, &completion_deps, decrement_module, ) diff --git a/src/scheduler/spawn/completion.rs b/src/scheduler/spawn/completion.rs index 54c0e33..a3d1f8c 100644 --- a/src/scheduler/spawn/completion.rs +++ b/src/scheduler/spawn/completion.rs @@ -30,6 +30,7 @@ pub(crate) async fn handle_success( task: &TaskRecord, phase: ExecutionPhase, metrics: &IoBudget, + memo: Option>, deps: &CompletionDeps, decrement_module: impl FnOnce(), ) { @@ -46,7 +47,7 @@ pub(crate) async fn handle_success( { match deps.store.active_children_count(task_id).await { Ok(count) if count > 0 => { - if let Err(e) = deps.store.set_waiting(task_id).await { + if let Err(e) = deps.store.set_waiting(task_id, memo.as_deref()).await { tracing::error!(task_id, error = %e, "failed to set task to waiting"); } decrement_module(); diff --git a/src/scheduler/tests.rs b/src/scheduler/tests.rs index 1834096..0a195f2 100644 --- a/src/scheduler/tests.rs +++ b/src/scheduler/tests.rs @@ -613,6 +613,7 @@ impl TypedExecutor for FinalizeTrackingExecutor { async fn finalize<'a>( &'a self, _payload: ParentTask, + _memo: (), _ctx: DomainTaskContext<'a, ParentDomain>, ) -> Result<(), TaskError> { self.finalized diff --git a/src/store/hierarchy.rs b/src/store/hierarchy.rs index b6d1815..6f60576 100644 --- a/src/store/hierarchy.rs +++ b/src/store/hierarchy.rs @@ -10,13 +10,15 @@ use super::{StoreError, TaskStore}; impl TaskStore { // ── Hierarchy ─────────────────────────────────────────────────── - /// Transition a running parent task to `waiting` status. + /// Transition a running parent task to `waiting` status, optionally + /// persisting a memo blob from `execute()`. /// /// Called after the parent's executor returns when it has spawned children. - pub async fn set_waiting(&self, id: i64) -> Result<(), StoreError> { + pub async fn set_waiting(&self, id: i64, memo: Option<&[u8]>) -> Result<(), StoreError> { sqlx::query( - "UPDATE tasks SET status = 'waiting', started_at = NULL WHERE id = ? AND status = 'running'", + "UPDATE tasks SET status = 'waiting', started_at = NULL, memo = ? WHERE id = ? AND status = 'running'", ) + .bind(memo) .bind(id) .execute(&self.pool) .await?; @@ -242,7 +244,7 @@ mod tests { store.submit(&sub).await.unwrap(); let task = store.pop_next().await.unwrap().unwrap(); - store.set_waiting(task.id).await.unwrap(); + store.set_waiting(task.id, None).await.unwrap(); let t = store.task_by_id(task.id).await.unwrap().unwrap(); assert_eq!(t.status, TaskStatus::Waiting); @@ -259,7 +261,7 @@ mod tests { let parent_sub = make_submission("parent", Priority::NORMAL); let parent_id = store.submit(&parent_sub).await.unwrap().id().unwrap(); store.pop_next().await.unwrap(); - store.set_waiting(parent_id).await.unwrap(); + store.set_waiting(parent_id, None).await.unwrap(); let mut child_sub = make_submission("child", Priority::NORMAL); child_sub.parent_id = Some(parent_id); @@ -281,7 +283,7 @@ mod tests { let parent_sub = make_submission("parent", Priority::NORMAL); let parent_id = store.submit(&parent_sub).await.unwrap().id().unwrap(); store.pop_next().await.unwrap(); - store.set_waiting(parent_id).await.unwrap(); + store.set_waiting(parent_id, None).await.unwrap(); for i in 0..2 { let mut sub = make_submission(&format!("child-{i}"), Priority::NORMAL); @@ -305,7 +307,7 @@ mod tests { let parent_sub = make_submission("parent", Priority::NORMAL); let parent_id = store.submit(&parent_sub).await.unwrap().id().unwrap(); store.pop_next().await.unwrap(); - store.set_waiting(parent_id).await.unwrap(); + store.set_waiting(parent_id, None).await.unwrap(); let mut child_sub = make_submission("child", Priority::NORMAL); child_sub.parent_id = Some(parent_id); @@ -402,7 +404,7 @@ mod tests { let sub = make_submission("fin", Priority::NORMAL); store.submit(&sub).await.unwrap(); let task = store.pop_next().await.unwrap().unwrap(); - store.set_waiting(task.id).await.unwrap(); + store.set_waiting(task.id, None).await.unwrap(); store.set_running_for_finalize(task.id).await.unwrap(); let t = store.task_by_id(task.id).await.unwrap().unwrap(); @@ -490,7 +492,7 @@ mod tests { let parent_sub = make_submission("parent", Priority::NORMAL); let parent_id = store.submit(&parent_sub).await.unwrap().id().unwrap(); store.pop_next().await.unwrap(); - store.set_waiting(parent_id).await.unwrap(); + store.set_waiting(parent_id, None).await.unwrap(); let mut child_sub = make_submission("child", Priority::NORMAL); child_sub.parent_id = Some(parent_id); diff --git a/src/store/lifecycle/mod.rs b/src/store/lifecycle/mod.rs index baf0197..2058578 100644 --- a/src/store/lifecycle/mod.rs +++ b/src/store/lifecycle/mod.rs @@ -73,8 +73,8 @@ pub(crate) async fn insert_history( expected_read_bytes, expected_write_bytes, expected_net_rx_bytes, expected_net_tx_bytes, actual_read_bytes, actual_write_bytes, actual_net_rx_bytes, actual_net_tx_bytes, retry_count, last_error, created_at, started_at, duration_ms, parent_id, fail_fast, group_key, - ttl_seconds, ttl_from, expires_at, run_after, max_retries) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ttl_seconds, ttl_from, expires_at, run_after, max_retries, memo) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(&task.task_type) .bind(&task.key) @@ -112,6 +112,7 @@ pub(crate) async fn insert_history( .map(|dt| dt.format("%Y-%m-%d %H:%M:%S").to_string()), ) .bind(task.max_retries) + .bind(&task.memo) .execute(&mut **conn) .await?; diff --git a/src/store/mod.rs b/src/store/mod.rs index 75d250a..8cb0439 100644 --- a/src/store/mod.rs +++ b/src/store/mod.rs @@ -273,6 +273,8 @@ impl TaskStore { include_str!("../../migrations/008_retry_backoff.sql"), ) .await?; + Self::run_alter_migration(&self.pool, include_str!("../../migrations/009_memo.sql")) + .await?; Ok(()) } diff --git a/src/store/row_mapping.rs b/src/store/row_mapping.rs index e12bd7d..bdb3833 100644 --- a/src/store/row_mapping.rs +++ b/src/store/row_mapping.rs @@ -119,6 +119,7 @@ pub(crate) fn row_to_task_record(row: &sqlx::sqlite::SqliteRow) -> TaskRecord { // Tags are populated separately from the task_tags table. tags: std::collections::HashMap::new(), max_retries: row.get("max_retries"), + memo: row.get("memo"), } } @@ -178,6 +179,7 @@ pub(crate) fn row_to_history_record(row: &sqlx::sqlite::SqliteRow) -> TaskHistor // Tags are populated separately from the task_history_tags table. tags: std::collections::HashMap::new(), max_retries: row.get("max_retries"), + memo: row.get("memo"), } } diff --git a/src/task/mod.rs b/src/task/mod.rs index 7edc378..1599a3f 100644 --- a/src/task/mod.rs +++ b/src/task/mod.rs @@ -243,6 +243,9 @@ pub struct TaskRecord { /// with pre-migration tasks). Resolved at submit time from: per-type /// retry policy → global `SchedulerConfig::max_retries`. pub max_retries: Option, + /// Serialized memo from `execute()`, delivered to `finalize()`. + /// `None` when no memo was produced (e.g. `Memo = ()`). + pub memo: Option>, } impl TaskRecord { @@ -319,6 +322,8 @@ pub struct TaskHistoryRecord { pub tags: HashMap, /// Per-task retry limit. `None` means use global default. pub max_retries: Option, + /// Serialized memo from `execute()`, preserved for debugging/observability. + pub memo: Option>, } /// IO budget for a task: expected or actual disk and network IO bytes. diff --git a/src/task/tests.rs b/src/task/tests.rs index b82d39b..05c83d4 100644 --- a/src/task/tests.rs +++ b/src/task/tests.rs @@ -254,6 +254,7 @@ fn event_header_includes_tags() { dependencies: Vec::new(), on_dependency_failure: super::submission::DependencyFailurePolicy::Cancel, max_retries: None, + memo: None, }; record.tags.insert("env".into(), "prod".into()); record.tags.insert("owner".into(), "alice".into()); diff --git a/tests/integration.rs b/tests/integration.rs index 824a3bb..d86544c 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -18,6 +18,8 @@ mod common; mod cross_module; #[path = "integration/dependencies.rs"] mod dependencies; +#[path = "integration/memo.rs"] +mod memo; #[path = "integration/module_features.rs"] mod module_features; #[path = "integration/modules.rs"] diff --git a/tests/integration/common.rs b/tests/integration/common.rs index 37d1590..feae5a5 100644 --- a/tests/integration/common.rs +++ b/tests/integration/common.rs @@ -306,6 +306,7 @@ where async fn finalize<'a>( &'a self, _payload: T, + _memo: (), _ctx: DomainTaskContext<'a, T::Domain>, ) -> Result<(), TaskError> { self.finalized.store(true, Ordering::SeqCst); diff --git a/tests/integration/memo.rs b/tests/integration/memo.rs new file mode 100644 index 0000000..5b4cbdc --- /dev/null +++ b/tests/integration/memo.rs @@ -0,0 +1,519 @@ +//! Integration tests for the execute-to-finalize memo feature. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use serde::{Deserialize, Serialize}; +use taskmill::{ + Domain, DomainKey, DomainTaskContext, Scheduler, SchedulerEvent, TaskError, TaskStore, + TypedExecutor, TypedTask, +}; +use tokio_util::sync::CancellationToken; + +use super::common::*; + +// ── Domain and task types ─────────────────────────────────────────── + +struct MemoDomain; +impl DomainKey for MemoDomain { + const NAME: &'static str = "memo"; +} + +#[derive(Debug, Default, Clone, Serialize, Deserialize)] +struct MemoParent; +impl TypedTask for MemoParent { + type Domain = MemoDomain; + const TASK_TYPE: &'static str = "parent"; +} + +#[derive(Debug, Default, Clone, Serialize, Deserialize)] +struct MemoChild; +impl TypedTask for MemoChild { + type Domain = MemoDomain; + const TASK_TYPE: &'static str = "child"; +} + +#[derive(Debug, Default, Clone, Serialize, Deserialize)] +struct MemoLeaf; +impl TypedTask for MemoLeaf { + type Domain = MemoDomain; + const TASK_TYPE: &'static str = "leaf"; +} + +// ── Memo type ─────────────────────────────────────────────────────── + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +struct ScanMemo { + scan_start_ns: i64, + batch_id: u64, +} + +// ── Executors ─────────────────────────────────────────────────────── + +/// Executor that returns a ScanMemo from execute() and verifies it in finalize(). +struct MemoRoundTripExecutor { + finalized: Arc, + memo_matched: Arc, +} + +impl TypedExecutor for MemoRoundTripExecutor { + async fn execute<'a>( + &'a self, + _payload: MemoParent, + ctx: DomainTaskContext<'a, MemoDomain>, + ) -> Result { + ctx.spawn_child_with(MemoChild) + .key("memo-child-0") + .await + .map_err(|e| TaskError::new(e.to_string()))?; + + Ok(ScanMemo { + scan_start_ns: 42_000_000, + batch_id: 99, + }) + } + + async fn finalize<'a>( + &'a self, + _payload: MemoParent, + memo: ScanMemo, + _ctx: DomainTaskContext<'a, MemoDomain>, + ) -> Result<(), TaskError> { + self.finalized.store(true, Ordering::SeqCst); + if memo.scan_start_ns == 42_000_000 && memo.batch_id == 99 { + self.memo_matched.store(true, Ordering::SeqCst); + } + Ok(()) + } +} + +/// Executor without memo — verifies backward compatibility. +struct NoMemoExecutor { + finalized: Arc, +} + +impl TypedExecutor for NoMemoExecutor { + async fn execute<'a>( + &'a self, + _payload: MemoParent, + ctx: DomainTaskContext<'a, MemoDomain>, + ) -> Result<(), TaskError> { + ctx.spawn_child_with(MemoChild) + .key("no-memo-child-0") + .await + .map_err(|e| TaskError::new(e.to_string()))?; + Ok(()) + } + + async fn finalize<'a>( + &'a self, + _payload: MemoParent, + _memo: (), + _ctx: DomainTaskContext<'a, MemoDomain>, + ) -> Result<(), TaskError> { + self.finalized.store(true, Ordering::SeqCst); + Ok(()) + } +} + +/// Executor that produces a memo but spawns no children (leaf task). +struct LeafWithMemoExecutor; + +impl TypedExecutor for LeafWithMemoExecutor { + async fn execute<'a>( + &'a self, + _payload: MemoLeaf, + _ctx: DomainTaskContext<'a, MemoDomain>, + ) -> Result { + Ok(ScanMemo { + scan_start_ns: 1, + batch_id: 2, + }) + } +} + +/// Executor whose memo fails to serialize. +struct BadMemoExecutor; + +#[derive(Debug, Clone, Deserialize)] +struct BadMemo; + +impl Serialize for BadMemo { + fn serialize(&self, _s: S) -> Result { + Err(serde::ser::Error::custom( + "intentional serialization failure", + )) + } +} + +impl TypedExecutor for BadMemoExecutor { + async fn execute<'a>( + &'a self, + _payload: MemoLeaf, + _ctx: DomainTaskContext<'a, MemoDomain>, + ) -> Result { + Ok(BadMemo) + } +} + +// ── Helper ────────────────────────────────────────────────────────── + +fn start_run_loop(sched: &Scheduler) -> (tokio::task::JoinHandle<()>, CancellationToken) { + let token = CancellationToken::new(); + let sched_clone = sched.clone(); + let token_clone = token.clone(); + let handle = tokio::spawn(async move { + sched_clone.run(token_clone).await; + }); + (handle, token) +} + +// ── Tests ─────────────────────────────────────────────────────────── + +/// 1. Basic round-trip: execute() returns ScanMemo, finalize() receives it. +#[tokio::test] +async fn memo_round_trip() { + let finalized = Arc::new(AtomicBool::new(false)); + let memo_matched = Arc::new(AtomicBool::new(false)); + + let sched = Scheduler::builder() + .store(TaskStore::open_memory().await.unwrap()) + .domain( + Domain::::new() + .task_memo(MemoRoundTripExecutor { + finalized: finalized.clone(), + memo_matched: memo_matched.clone(), + }) + .task::(NoopExecutor), + ) + .build() + .await + .unwrap(); + + let handle = sched.domain::(); + let mut rx = sched.subscribe(); + + handle.submit(MemoParent).await.unwrap(); + + let (run_handle, token) = start_run_loop(&sched); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + let completed = wait_for_event( + &mut rx, + deadline, + |evt| matches!(evt, SchedulerEvent::Completed(h) if h.task_type == "memo::parent"), + ) + .await; + + token.cancel(); + let _ = run_handle.await; + + assert!(completed.is_some(), "parent should complete"); + assert!( + finalized.load(Ordering::SeqCst), + "finalize should be called" + ); + assert!( + memo_matched.load(Ordering::SeqCst), + "memo values should match in finalize" + ); +} + +/// 2. Executor with Memo = () works unchanged — no memo written to DB. +#[tokio::test] +async fn unit_memo_works_unchanged() { + let finalized = Arc::new(AtomicBool::new(false)); + + let sched = Scheduler::builder() + .store(TaskStore::open_memory().await.unwrap()) + .domain( + Domain::::new() + .task::(NoMemoExecutor { + finalized: finalized.clone(), + }) + .task::(NoopExecutor), + ) + .build() + .await + .unwrap(); + + let handle = sched.domain::(); + let mut rx = sched.subscribe(); + + handle.submit(MemoParent).await.unwrap(); + + let (run_handle, token) = start_run_loop(&sched); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + let completed = wait_for_event( + &mut rx, + deadline, + |evt| matches!(evt, SchedulerEvent::Completed(h) if h.task_type == "memo::parent"), + ) + .await; + + assert!(completed.is_some(), "parent should complete"); + assert!( + finalized.load(Ordering::SeqCst), + "finalize should be called for unit-memo executor" + ); + + // Verify no memo was persisted in history (query before shutdown closes pool). + let history = sched.store().history(10, 0).await.unwrap(); + let parent_hist = history + .iter() + .find(|h| h.task_type == "memo::parent") + .expect("parent should be in history"); + assert!( + parent_hist.memo.is_none(), + "unit memo should not be persisted" + ); + + token.cancel(); + let _ = run_handle.await; +} + +/// 3. Memo is preserved in history after task completes. +#[tokio::test] +async fn memo_preserved_in_history() { + let finalized = Arc::new(AtomicBool::new(false)); + let memo_matched = Arc::new(AtomicBool::new(false)); + + let sched = Scheduler::builder() + .store(TaskStore::open_memory().await.unwrap()) + .domain( + Domain::::new() + .task_memo(MemoRoundTripExecutor { + finalized: finalized.clone(), + memo_matched: memo_matched.clone(), + }) + .task::(NoopExecutor), + ) + .build() + .await + .unwrap(); + + let handle = sched.domain::(); + let mut rx = sched.subscribe(); + + handle.submit(MemoParent).await.unwrap(); + + let (run_handle, token) = start_run_loop(&sched); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + let _ = wait_for_event( + &mut rx, + deadline, + |evt| matches!(evt, SchedulerEvent::Completed(h) if h.task_type == "memo::parent"), + ) + .await; + + // Query before shutdown closes the pool. + let history = sched.store().history(10, 0).await.unwrap(); + let parent_hist = history + .iter() + .find(|h| h.task_type == "memo::parent") + .expect("parent should be in history"); + + assert!( + parent_hist.memo.is_some(), + "memo should be preserved in history" + ); + let memo: ScanMemo = serde_json::from_slice(parent_hist.memo.as_ref().unwrap()).unwrap(); + assert_eq!(memo.scan_start_ns, 42_000_000); + assert_eq!(memo.batch_id, 99); + + token.cancel(); + let _ = run_handle.await; +} + +/// 4. Leaf task (no children, no finalize): execute returns memo but task +/// completes normally — memo is not written since there's no set_waiting. +#[tokio::test] +async fn leaf_task_memo_not_persisted() { + let sched = Scheduler::builder() + .store(TaskStore::open_memory().await.unwrap()) + .domain(Domain::::new().task_memo(LeafWithMemoExecutor)) + .build() + .await + .unwrap(); + + let handle = sched.domain::(); + let mut rx = sched.subscribe(); + + handle.submit(MemoLeaf).await.unwrap(); + + let (run_handle, token) = start_run_loop(&sched); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + let _ = wait_for_event( + &mut rx, + deadline, + |evt| matches!(evt, SchedulerEvent::Completed(h) if h.task_type == "memo::leaf"), + ) + .await; + + let history = sched.store().history(10, 0).await.unwrap(); + let leaf_hist = history + .iter() + .find(|h| h.task_type == "memo::leaf") + .expect("leaf should be in history"); + assert!( + leaf_hist.memo.is_none(), + "leaf task memo should not be persisted (no waiting transition)" + ); + + token.cancel(); + let _ = run_handle.await; +} + +/// 5. Serialization failure: executor returns a type that fails to serialize → +/// execute returns TaskError::permanent, task is not left in a broken state. +#[tokio::test] +async fn serialization_failure_produces_permanent_error() { + let sched = Scheduler::builder() + .store(TaskStore::open_memory().await.unwrap()) + .domain(Domain::::new().task_memo(BadMemoExecutor)) + .build() + .await + .unwrap(); + + let handle = sched.domain::(); + let mut rx = sched.subscribe(); + + handle.submit(MemoLeaf).await.unwrap(); + + let (run_handle, token) = start_run_loop(&sched); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + let failed = wait_for_event(&mut rx, deadline, |evt| { + matches!( + evt, + SchedulerEvent::Failed { + header, + will_retry: false, + .. + } if header.task_type == "memo::leaf" + ) + }) + .await; + + token.cancel(); + let _ = run_handle.await; + + assert!( + failed.is_some(), + "task should fail with serialization error" + ); + + if let Some(SchedulerEvent::Failed { error, .. }) = failed { + assert!( + error.contains("memo serialization"), + "error should mention memo serialization: {error}" + ); + } +} + +/// 6. Memo survives restart: persist memo → stop scheduler → reopen → +/// children complete → finalize receives correct memo. +#[tokio::test] +async fn memo_survives_restart() { + let dir = std::env::temp_dir().join(format!("taskmill-memo-restart-{}", std::process::id())); + let _ = std::fs::create_dir_all(&dir); + let db_path = dir.join("test.db"); + let db_str = db_path.to_str().unwrap(); + + let finalized = Arc::new(AtomicBool::new(false)); + let memo_matched = Arc::new(AtomicBool::new(false)); + + // Phase 1: submit parent, let it execute (spawns child), enter waiting. + { + let sched = Scheduler::builder() + .store(TaskStore::open(db_str).await.unwrap()) + .domain( + Domain::::new() + .task_memo(MemoRoundTripExecutor { + finalized: finalized.clone(), + memo_matched: memo_matched.clone(), + }) + .task::(DelayExecutor(Duration::from_secs(60))), + ) + .build() + .await + .unwrap(); + + let handle = sched.domain::(); + let mut rx = sched.subscribe(); + + handle.submit(MemoParent).await.unwrap(); + + let (run_handle, token) = start_run_loop(&sched); + + // Wait for parent to enter waiting. + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + let _ = wait_for_event(&mut rx, deadline, |evt| { + matches!(evt, SchedulerEvent::Waiting { .. }) + }) + .await; + + // Verify memo is persisted on the task record. + let tasks = sched.store().waiting_tasks().await.unwrap(); + assert_eq!(tasks.len(), 1, "parent should be in waiting state"); + assert!(tasks[0].memo.is_some(), "memo should be persisted"); + let memo: ScanMemo = serde_json::from_slice(tasks[0].memo.as_ref().unwrap()).unwrap(); + assert_eq!(memo.scan_start_ns, 42_000_000); + assert_eq!(memo.batch_id, 99); + + // Stop the scheduler (simulates restart). + token.cancel(); + let _ = run_handle.await; + sched.store().close().await; + } + + // Phase 2: reopen from same DB — child re-runs and completes, parent finalizes. + { + let sched = Scheduler::builder() + .store(TaskStore::open(db_str).await.unwrap()) + .domain( + Domain::::new() + .task_memo(MemoRoundTripExecutor { + finalized: finalized.clone(), + memo_matched: memo_matched.clone(), + }) + .task::(NoopExecutor), + ) + .build() + .await + .unwrap(); + + let mut rx = sched.subscribe(); + + let (run_handle, token) = start_run_loop(&sched); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + let _ = wait_for_event( + &mut rx, + deadline, + |evt| matches!(evt, SchedulerEvent::Completed(h) if h.task_type == "memo::parent"), + ) + .await; + + token.cancel(); + let _ = run_handle.await; + + assert!( + finalized.load(Ordering::SeqCst), + "finalize should run after restart" + ); + assert!( + memo_matched.load(Ordering::SeqCst), + "memo should survive restart and match in finalize" + ); + + sched.store().close().await; + } + + // Cleanup. + let _ = std::fs::remove_dir_all(&dir); +}