swaits/dagx
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: make DagRunner interface single-threaded

Browse Source
This commit is contained in:
TechnoPorg committed 2026-02-17 21:29:14 +01:00
1 parent d83a815ed0
commit e4fc0b0d44
25 files changed
+416 -617
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/builder.rs
+140 -16
View File
@@ -3,10 +3,11 @@
//! TaskBuilder uses compile-time type checking to ensure dependencies are wired correctly.
//! The type parameter tracks whether dependencies have been specified yet.
use std::marker::PhantomData;
use crate::deps::DepsTuple;
use crate::runner::DagRunner;
use crate::task::Task;
use crate::types::{NodeId, TaskHandle};
#[cfg(feature = "tracing")]
use tracing::debug;
@@ -57,17 +58,26 @@ use tracing::debug;
/// let b = b.depends_on(a);
/// // Now b is a TaskHandle<i32>
///
/// dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
/// assert_eq!(dag.get(b).unwrap(), 20);
///let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
/// assert_eq!(output.get(b).unwrap(), 20);
/// # };
/// ```
pub struct TaskBuilder<'a, Tk: Task, Deps> {
#[must_use]
pub struct TaskBuilder<'a, Input, Tk>
where
Tk: Task<Input>,
Input: Send + Sync + 'static,
{
pub(crate) id: NodeId,
pub(crate) dag: &'a DagRunner,
pub(crate) _phantom: std::marker::PhantomData<(Tk, Deps)>,
pub(crate) dag: &'a mut DagRunner,
pub(crate) _phantom: PhantomData<(Tk, Input)>,
}
impl<'a, Tk: Task, Deps> TaskBuilder<'a, Tk, Deps> {
impl<'a, Input, Tk> TaskBuilder<'a, Input, Tk>
where
Tk: Task<Input>,
Input: Send + Sync + 'static,
{
/// Provide all dependencies exactly once as a tuple.
///
/// The dependencies must match the task's `Input` type exactly:
@@ -107,7 +117,7 @@ impl<'a, Tk: Task, Deps> TaskBuilder<'a, Tk, Deps> {
/// }
///
/// # async {
/// let dag = DagRunner::new();
/// let mut dag = DagRunner::new();
///
/// let x = dag.add_task(Value(2)).into();
/// let y = dag.add_task(Value(3));
@@ -118,13 +128,13 @@ impl<'a, Tk: Task, Deps> TaskBuilder<'a, Tk, Deps> {
/// // Multiple dependencies: tuple form
/// let sum = dag.add_task(Add).depends_on((&x, y));
///
/// dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
///let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
/// # };
/// ```
#[allow(private_bounds)]
pub fn depends_on<D>(self, deps: D) -> TaskHandle<Tk::Output>
where
D: DepsTuple<Tk::Input>,
D: DepsTuple<Input>,
{
// Register dependencies in the DAG
let dep_ids = deps.to_node_ids();
@@ -137,27 +147,141 @@ impl<'a, Tk: Task, Deps> TaskBuilder<'a, Tk, Deps> {
"wiring task dependencies"
);
let mut edges = self.dag.edges.lock();
let mut dependents = self.dag.dependents.lock();
for &dep_id in &dep_ids {
// Add edge from this node to dependency
if let Some(node_edges) = edges.get_mut(&self.id) {
if let Some(node_edges) = self.dag.edges.get_mut(&self.id) {
node_edges.push(dep_id);
}
// Add this node as dependent of the dependency
if let Some(node_dependents) = dependents.get_mut(&dep_id) {
if let Some(node_dependents) = self.dag.dependents.get_mut(&dep_id) {
node_dependents.push(self.id);
}
}
TaskHandle {
id: self.id,
_phantom: std::marker::PhantomData,
_phantom: PhantomData,
}
}
}
/// Opaque node identifier
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub struct NodeId(pub u32);
/// Opaque, typed token for a node's output.
///
/// A `TaskHandle<T>` provides compile-time type-safe access to a task's output.
/// You can:
/// 1. Pass it to [`crate::TaskBuilder::depends_on`] to wire up dependencies
/// 2. Use it with [`crate::DagOutput::get`] to retrieve the output after [`crate::DagRunner::run`]
///
/// Handles are cheap to clone and copy.
///
/// # Examples
///
/// ```no_run
/// # use dagx::{task, DagRunner, Task};
/// #
/// # struct LoadValue { value: i32 }
/// # impl LoadValue { pub fn new(v: i32) -> Self { Self { value: v } } }
/// # #[task]
/// # impl LoadValue {
/// # async fn run(&mut self) -> i32 { self.value }
/// # }
/// # async {
/// let mut dag = DagRunner::new();
/// let node = dag.add_task(LoadValue::new(42));
///
///let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
///
/// assert_eq!(output.get(node).unwrap(), 42);
/// # };
/// ```
pub struct TaskHandle<T> {
pub(crate) id: NodeId,
pub(crate) _phantom: PhantomData<fn() -> T>,
}
impl<T> Clone for TaskHandle<T> {
fn clone(&self) -> Self {
*self
}
}
impl<T> Copy for TaskHandle<T> {}
// TaskHandle can be converted from &TaskHandle (for .get() calls)
impl<T> From<&TaskHandle<T>> for TaskHandle<T> {
fn from(handle: &TaskHandle<T>) -> Self {
*handle
}
}
/// Takes a task and converts it to either a TaskBuilder or a TaskHandle,
/// depending on whether it has inputs or not.
///
/// This is useful to enforce at compile time that a task with unit input never has
/// depends_on called, and that it can be used directly as a dependency
/// without converting it manually to a TaskHandle.
pub trait TaskWire<Input>: Task<Input> + Sync + 'static
where
Input: Send + Sync + 'static,
{
type Retval<'dag>;
fn new_from_dag<'dag>(id: NodeId, dag: &'dag mut DagRunner) -> Self::Retval<'dag>;
}
impl<Tk> TaskWire<()> for Tk
where
Tk: Task<()> + Sync + 'static,
{
type Retval<'dag> = TaskHandle<Tk::Output>;
fn new_from_dag(id: NodeId, _dag: &mut DagRunner) -> Self::Retval<'static> {
Self::Retval {
id,
_phantom: PhantomData,
}
}
}
/// Macro to implement TaskWire for different tuple sizes.
macro_rules! impl_wire_tuple {
($($T:ident),+) => {
// For tuples of any combination of &TaskHandle, TaskHandle, and TaskBuilder
impl<Tk, $($T: Send + Sync + 'static),+> TaskWire<($($T,)+)> for Tk
where
Tk: Task<($($T,)+)> + Sync + 'static
{
type Retval<'dag> = TaskBuilder<'dag, ($($T,)+), Tk>;
fn new_from_dag<'dag>(id: NodeId, dag: &'dag mut DagRunner) -> Self::Retval<'dag> {
Self::Retval {
id,
dag,
_phantom: PhantomData,
}
}
}
};
}
// Generate DepsTuple implementations for tuples of size 2-8.
// Supporting up to 8 elements covers the vast majority of use cases.
impl_wire_tuple!(T1);
impl_wire_tuple!(T1, T2);
impl_wire_tuple!(T1, T2, T3);
impl_wire_tuple!(T1, T2, T3, T4);
impl_wire_tuple!(T1, T2, T3, T4, T5);
impl_wire_tuple!(T1, T2, T3, T4, T5, T6);
impl_wire_tuple!(T1, T2, T3, T4, T5, T6, T7);
impl_wire_tuple!(T1, T2, T3, T4, T5, T6, T7, T8);
#[cfg(test)]
mod tests;
#[cfg(test)]
mod coverage_tests;
src/types/coverage_tests.rs → src/builder/coverage_tests.rs
+7 -6
View File
@@ -25,19 +25,20 @@ async fn test_unit_struct_with_state() {
}
}
let dag = DagRunner::new();
let mut dag = DagRunner::new();
// Create task with initial state
let builder = dag.add_task(StatefulTask::new(21));
// Explicit conversion - triggers line 71
let handle: TaskHandle<_> = builder.into();
let handle = builder;
dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
let mut output = dag
.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
.await
.unwrap();
assert_eq!(dag.get(handle).unwrap(), 42);
assert_eq!(output.get(handle).unwrap(), 42);
}
#[test]
@@ -51,12 +52,12 @@ fn test_taskbuilder_conversion_without_running() {
}
}
let dag = DagRunner::new();
let mut dag = DagRunner::new();
let builder = dag.add_task(TestTask);
let builder_id = builder.id;
// Convert without running the DAG
let handle: TaskHandle<String> = builder.into();
let handle: TaskHandle<String> = builder;
// Verify the ID was transferred correctly
assert_eq!(handle.id, builder_id);
src/builder/tests.rs
+10 -11
View File
@@ -1,7 +1,6 @@
//! Unit tests for builder module
use crate::runner::DagRunner;
use crate::types::TaskHandle;
// Initialize tracing subscriber for tests (idempotent)
#[cfg(feature = "tracing")]
@@ -47,11 +46,11 @@ impl TestTaskWithInput {
#[test]
fn test_task_builder_from_conversion() {
// Test line 165 in builder.rs - the From implementation
let dag = DagRunner::new();
let mut dag = DagRunner::new();
let builder = dag.add_task(TestTask { value: 42 });
// Convert builder to TaskHandle
let handle: TaskHandle<_> = builder.into();
let handle = builder;
// Check that the ID is preserved
assert_eq!(handle.id.0, 0); // First task should have ID 0
@@ -59,17 +58,17 @@ fn test_task_builder_from_conversion() {
#[test]
fn test_task_builder_chain() {
let dag = DagRunner::new();
let mut dag = DagRunner::new();
// Test chaining multiple tasks
let t1 = dag.add_task(TestTask { value: 1 }).into();
let t1 = dag.add_task(TestTask { value: 1 });
let t2 = dag.add_task(TestTaskWithInput).depends_on(t1);
let t3 = dag.add_task(TestTaskWithInput).depends_on(t2);
// All should return TaskHandles
let _: TaskHandle<_> = t1;
let _: TaskHandle<_> = t2;
let _: TaskHandle<_> = t3;
let _ = t1;
let _ = t2;
let _ = t3;
}
#[test]
@@ -84,7 +83,7 @@ fn test_multiple_dependencies() {
}
}
let dag = DagRunner::new();
let mut dag = DagRunner::new();
let a = dag.add_task(TestTask { value: 10 });
let b = dag.add_task(TestTask { value: 20 });
@@ -93,12 +92,12 @@ fn test_multiple_dependencies() {
let sum = dag.add_task(AddTask).depends_on((a, b));
// Should return a TaskHandle
let _: TaskHandle<_> = sum;
let _ = sum;
}
#[test]
fn test_task_builder_stores_correct_id() {
let dag = DagRunner::new();
let mut dag = DagRunner::new();
let builder1 = dag.add_task(TestTask { value: 1 });
let builder2 = dag.add_task(TestTask { value: 2 });
src/deps.rs
+18 -46
View File
@@ -7,36 +7,7 @@
//! This trait is internal and not meant for external implementation. Users should only
//! use the provided implementations.
use crate::builder::TaskBuilder;
use crate::task::Task;
use crate::types::{NodeId, TaskHandle};
trait IntoNodeRef {
type Output;
fn into_node_ref(self) -> NodeId;
}
impl<T> IntoNodeRef for &TaskHandle<T> {
type Output = T;
fn into_node_ref(self) -> NodeId {
self.id
}
}
impl<T> IntoNodeRef for TaskHandle<T> {
type Output = T;
fn into_node_ref(self) -> NodeId {
self.id
}
}
impl<'runner, Tk: Task, Deps> IntoNodeRef for TaskBuilder<'runner, Tk, Deps> {
type Output = Tk::Output;
fn into_node_ref(self) -> NodeId {
self.id
}
}
use crate::builder::{NodeId, TaskHandle};
pub(crate) trait DepsTuple<Input> {
fn to_node_ids(self) -> Vec<NodeId>;
@@ -50,21 +21,21 @@ impl DepsTuple<()> for () {
}
// Implementation for single dependency (all forms)
impl<T> DepsTuple<(T::Output,)> for T
impl<T, O> DepsTuple<(O,)> for T
where
T: IntoNodeRef,
T: Into<TaskHandle<O>>,
{
fn to_node_ids(self) -> Vec<NodeId> {
vec![self.into_node_ref()]
vec![self.into().id]
}
}
impl<T> DepsTuple<T::Output> for (T,)
impl<T, O> DepsTuple<O> for (T,)
where
T: IntoNodeRef,
T: Into<TaskHandle<O>>,
{
fn to_node_ids(self) -> Vec<NodeId> {
vec![self.0.into_node_ref()]
vec![self.0.into().id]
}
}
@@ -82,13 +53,14 @@ where
/// task.depends_on((&dag.add_task(TaskA), &dag.add_task(TaskB)))
/// ```
macro_rules! impl_deps_tuple {
($($T:ident),+) => {
($($T:ident : $O:ident),+) => {
// For tuples of any combination of &TaskHandle, TaskHandle, and TaskBuilder
impl<$($T: IntoNodeRef),+> DepsTuple<($($T::Output,)+)> for ($($T,)+) {
impl<$($T: Into<TaskHandle<$O>>, $O),+> DepsTuple<($($O,)+)> for ($($T,)+)
{
#[allow(non_snake_case)]
fn to_node_ids(self) -> Vec<NodeId> {
let ($($T,)+) = self;
vec![$($T.into_node_ref(),)+]
vec![$($T.into().id,)+]
}
}
};
@@ -96,13 +68,13 @@ macro_rules! impl_deps_tuple {
// Generate DepsTuple implementations for tuples of size 2-8.
// Supporting up to 8 elements covers the vast majority of use cases.
impl_deps_tuple!(T1, T2);
impl_deps_tuple!(T1, T2, T3);
impl_deps_tuple!(T1, T2, T3, T4);
impl_deps_tuple!(T1, T2, T3, T4, T5);
impl_deps_tuple!(T1, T2, T3, T4, T5, T6);
impl_deps_tuple!(T1, T2, T3, T4, T5, T6, T7);
impl_deps_tuple!(T1, T2, T3, T4, T5, T6, T7, T8);
impl_deps_tuple!(T1:O1, T2:O2);
impl_deps_tuple!(T1:O1, T2:O2, T3:O3);
impl_deps_tuple!(T1:O1, T2:O2, T3:O3, T4:O4);
impl_deps_tuple!(T1:O1, T2:O2, T3:O3, T4:O4, T5:O5);
impl_deps_tuple!(T1:O1, T2:O2, T3:O3, T4:O4, T5:O5, T6:O6);
impl_deps_tuple!(T1:O1, T2:O2, T3:O3, T4:O4, T5:O5, T6:O6, T7:O7);
impl_deps_tuple!(T1:O1, T2:O2, T3:O3, T4:O4, T5:O5, T6:O6, T7:O7, T8:O8);
#[cfg(test)]
mod tests;
src/deps/tests.rs
+4 -4
View File
@@ -1,8 +1,8 @@
//! Unit tests for deps module
use crate::builder::{NodeId, TaskHandle};
use crate::deps::DepsTuple;
use crate::runner::DagRunner;
use crate::types::{NodeId, TaskHandle};
// Test task for dependency testing
struct TestTask {
@@ -78,7 +78,7 @@ fn test_deps_tuple_single_handle_tuple_owned() {
#[test]
fn test_deps_tuple_task_builder_ref() {
// Test lines 51-55 in deps.rs - &TaskBuilder
let dag = DagRunner::new();
let mut dag = DagRunner::new();
let builder = dag.add_task(TestTask { value: 42 });
let builder_id = builder.id;
@@ -92,7 +92,7 @@ fn test_deps_tuple_task_builder_ref() {
#[test]
fn test_deps_tuple_task_builder_tuple() {
// Test lines 57-61 in deps.rs - (&TaskBuilder,)
let dag = DagRunner::new();
let mut dag = DagRunner::new();
let builder = dag.add_task(TestTask { value: 100 });
let builder_id = builder.id;
@@ -134,7 +134,7 @@ fn test_deps_tuple_multiple_handles() {
#[test]
fn test_deps_tuple_multiple_builders() {
// Test macro-generated implementations for multiple builders
let dag = DagRunner::new();
let mut dag = DagRunner::new();
let builder1 = dag.add_task(TestTask { value: 10 });
let builder1_id = builder1.id;
let builder2 = dag.add_task(TestTask { value: 20 });
src/error.rs
-10
View File
@@ -17,8 +17,6 @@ pub enum DagError {
TaskPanicked { task_id: u32, panic_message: String },
/// Result not found for task
ResultNotFound { task_id: u32 },
/// Attempted to run the DAG while it's already running
ConcurrentExecution,
}
impl std::fmt::Display for DagError {
@@ -65,14 +63,6 @@ impl std::fmt::Display for DagError {
task_id
)
}
DagError::ConcurrentExecution => {
write!(
f,
"DAG is already running - concurrent execution not supported.\n\
\n\
Wait for the current execution to complete before starting another."
)
}
}
}
}
src/error/tests.rs
+3 -11
View File
@@ -51,17 +51,9 @@ fn test_dag_error_display_result_not_found() {
assert!(display.contains("Result not found"));
assert!(display.contains("task #7"));
assert!(display.contains("Call dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })"));
}
#[test]
fn test_dag_error_display_concurrent_execution() {
let err = DagError::ConcurrentExecution;
let display = format!("{}", err);
assert!(display.contains("already running"));
assert!(display.contains("concurrent execution not supported"));
assert!(display.contains("Wait for the current execution"));
assert!(display.contains(
"Calllet mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })"
));
}
#[test]
src/lib.rs
+36 -36
View File
@@ -42,18 +42,18 @@
//! }
//!
//! # async {
//! let dag = DagRunner::new();
//! let mut dag = DagRunner::new();
//!
//! // Add source tasks
//! let x = dag.add_task(Value(2));
//! let y: TaskHandle<_> = dag.add_task(Value(3)).into(); // Allows y to be reused as a dependency
//! let y = dag.add_task(Value(3)).into(); // Allows y to be reused as a dependency
//!
//! // Add task with dependencies
//! let sum = dag.add_task(Add).depends_on((x, &y));
//!
//! // Execute and retrieve results
//! dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//! assert_eq!(dag.get(sum).unwrap(), 5);
//!let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//! assert_eq!(output.get(sum).unwrap(), 5);
//! # };
//! ```
//!
@@ -212,7 +212,7 @@
//!
//! A [`TaskHandle<T>`] is a typed, opaque reference to a task's output. Use it to:
//! - Wire dependencies between tasks
//! - Retrieve results after execution with [`DagRunner::get`]
//! - Retrieve results after execution with [`DagOutput::get`]
//!
//! ## Dependency Patterns
//!
@@ -230,7 +230,7 @@
//! # async fn run(&self) -> i32 { self.0 }
//! # }
//! # async {
//! # let dag = DagRunner::new();
//! # let mut dag = DagRunner::new();
//! let source = dag.add_task(Value(42));
//! # };
//! ```
@@ -252,7 +252,7 @@
//! async fn run(input: &i32) -> i32 { input * 2 }
//! }
//! # async {
//! # let dag = DagRunner::new();
//! # let mut dag = DagRunner::new();
//! # let source = dag.add_task(Value(42));
//! let doubled = dag.add_task(Double).depends_on(source);
//! # };
@@ -275,7 +275,7 @@
//! async fn run(a: &i32, b: &i32) -> i32 { a + b }
//! }
//! # async {
//! # let dag = DagRunner::new();
//! # let mut dag = DagRunner::new();
//! # let x = dag.add_task(Value(2));
//! # let y = dag.add_task(Value(3));
//! let sum = dag.add_task(Add).depends_on((x, y));
@@ -317,16 +317,16 @@
//! }
//!
//! # async {
//! let dag = DagRunner::new();
//! let mut dag = DagRunner::new();
//!
//! let base = dag.add_task(Value(10)).into();
//! let plus1 = dag.add_task(Add(1)).depends_on(&base);
//! let times2 = dag.add_task(Scale(2)).depends_on(&base);
//!
//! dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//!let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//!
//! assert_eq!(dag.get(plus1).unwrap(), 11);
//! assert_eq!(dag.get(times2).unwrap(), 20);
//! assert_eq!(output.get(plus1).unwrap(), 11);
//! assert_eq!(output.get(times2).unwrap(), 20);
//! # };
//! ```
//!
@@ -373,16 +373,16 @@
//! }
//!
//! # async {
//! let dag = DagRunner::new();
//! let mut dag = DagRunner::new();
//!
//! let name = dag.add_task(Name("Alice".to_string()));
//! let age = dag.add_task(Age(30));
//! let active = dag.add_task(Active(true));
//! let result = dag.add_task(FormatUser).depends_on((name, age, active));
//!
//! dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//!let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//!
//! assert_eq!(dag.get(result).unwrap(), "User: Alice, Age: 30, Active: true");
//! assert_eq!(output.get(result).unwrap(), "User: Alice, Age: 30, Active: true");
//! # };
//! ```
//!
@@ -419,7 +419,7 @@
//! }
//!
//! # async {
//! let dag = DagRunner::new();
//! let mut dag = DagRunner::new();
//!
//! // Layer 1: Sources
//! let x = dag.add_task(Value(2));
@@ -433,9 +433,9 @@
//! // Layer 3: Final result
//! let total = dag.add_task(Add).depends_on((&sum_xy, &prod_yz)); // 5 + 15 = 20
//!
//! dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//!let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//!
//! assert_eq!(dag.get(total).unwrap(), 20);
//! assert_eq!(output.get(total).unwrap(), 20);
//! # };
//! ```
//!
@@ -470,14 +470,14 @@
//! }
//!
//! # async {
//! let dag = DagRunner::new();
//! let mut dag = DagRunner::new();
//!
//! let user = dag.add_task(CreateUser);
//! let formatted = dag.add_task(FormatUser).depends_on(user);
//!
//! dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//!let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//!
//! assert_eq!(dag.get(formatted).unwrap(), "Alice is 30 years old");
//! assert_eq!(output.get(formatted).unwrap(), "Alice is 30 years old");
//! # };
//! ```
//!
@@ -497,21 +497,21 @@
//! // With Tokio
//! #[tokio::main]
//! async fn main() {
//! let dag = DagRunner::new();
//! let mut dag = DagRunner::new();
//! // ... build and run DAG
//! }
//!
//! // With async-std
//! #[async_std::main]
//! async fn main() {
//! let dag = DagRunner::new();
//! let mut dag = DagRunner::new();
//! // ... build and run DAG
//! }
//!
//! // With smol
//! fn main() {
//! smol::block_on(async {
//! let dag = DagRunner::new();
//! let mut dag = DagRunner::new();
//! // ... build and run DAG
//! });
//! }
@@ -522,9 +522,9 @@
//! dagx uses [`DagResult<T>`] (an alias for `Result<T, DagError>`) for operations that
//! can fail:
//!
//! - [`DagRunner::run`] returns `DagResult<()>` and can fail if tasks panic or if multiple
//! - [`DagRunner::run`] returns `DagResult<DagOutput>` and can fail if tasks panic or if multiple
//! concurrent runs are attempted
//! - [`DagRunner::get`] returns `DagResult<T>` and can fail if the task hasn't executed or
//! - [`DagOutput::get`] returns `DagResult<T>` and can fail if the task hasn't executed or
//! the handle is invalid
//!
//! You can handle errors explicitly or use `.unwrap()` for simple cases:
@@ -538,14 +538,14 @@
//! # async fn run(&self) -> i32 { self.0 }
//! # }
//! # async {
//! # let dag = DagRunner::new();
//! # let mut dag = DagRunner::new();
//! # let node = dag.add_task(Value(42));
//! // Simple approach with .unwrap()
//! dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//! let result = dag.get(node).unwrap();
//!let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//! let result = output.get(node).unwrap();
//!
//! // Or handle errors explicitly
//! match dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await {
//! matchlet mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await {
//! Ok(_) => println!("DAG executed successfully"),
//! Err(e) => eprintln!("DAG execution failed: {}", e),
//! }
@@ -651,7 +651,7 @@
//! }
//!
//! # async {
//! let dag = DagRunner::new();
//! let mut dag = DagRunner::new();
//! let data = dag.add_task(FetchData).into();
//!
//! // All three tasks get efficient Arc-wrapped sharing automatically
@@ -659,7 +659,7 @@
//! let task2 = dag.add_task(ProcessData).depends_on(&data);
//! let task3 = dag.add_task(ProcessData).depends_on(&data);
//!
//! dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//!let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
//! # };
//! ```
//!
@@ -668,7 +668,7 @@
//! - Framework wraps it in `Arc<T>` internally
//! - For fan-out (1→N), Arc is cloned N times (just atomic pointer increments - O(1))
//! - Each downstream task receives `&T` after extracting from Arc
//! - When you call `dag.get()`, the Arc is cloned and inner value extracted
//! - When you call `output.get()`, the Arc is cloned and inner value extracted
//!
//! **Performance characteristics:**
//! - **Heap types** (Vec, String, HashMap): Arc overhead is negligible (~few ns)
@@ -771,16 +771,16 @@ mod builder;
mod deps;
mod error;
mod node;
mod output;
mod runner;
mod task;
mod types;
// Public re-exports
pub use builder::TaskBuilder;
pub use builder::{TaskBuilder, TaskHandle};
pub use error::{DagError, DagResult};
pub use output::DagOutput;
pub use runner::DagRunner;
pub use task::{Task, TaskInput};
pub use types::{Pending, TaskHandle};
// Re-export the procedural macro
pub use dagx_macros::task;
src/node.rs
+19 -6
View File
@@ -8,6 +8,7 @@
use std::any::Any;
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::Arc;
@@ -50,20 +51,32 @@ pub(crate) trait ExecutableNode: Send {
/// Each task executes exactly once, taking ownership and producing an output.
///
/// The output is also returned from execute_with_deps for central storage.
pub(crate) struct TypedNode<T: Task> {
pub(crate) struct TypedNode<Input, T>
where
Input: Send + Sync + 'static,
T: Task<Input> + 'static,
{
pub(crate) task: T,
phantom: PhantomData<Input>,
}
impl<T: Task + 'static> TypedNode<T> {
impl<Input, T> TypedNode<Input, T>
where
T: Task<Input>,
Input: Send + Sync + 'static,
{
pub(crate) fn new(task: T) -> Self {
Self { task }
Self {
task,
phantom: PhantomData,
}
}
}
impl<T: Task + 'static> ExecutableNode for TypedNode<T>
impl<Input, T> ExecutableNode for TypedNode<Input, T>
where
T::Input: 'static,
T::Output: 'static,
T: Task<Input>,
Input: Send + Sync + 'static,
{
fn execute_with_deps(
self: Box<Self>,
src/output.rs
+83
View File
@@ -0,0 +1,83 @@
use std::{any::Any, sync::Arc};
use crate::{builder::NodeId, runner::PassThroughHashMap, DagError, DagResult, TaskHandle};
pub struct DagOutput {
outputs: PassThroughHashMap<NodeId, Arc<dyn Any + Send + Sync + 'static>>,
}
impl DagOutput {
/// Retrieve a task's output after [`crate::DagRunner::run`].
///
/// # Behavior
///
/// All task outputs are stored after execution and can be retrieved via get().
/// Each task's output can only be taken once, after which it will return `DagError::ResultNotFound`.
///
/// # Errors
///
/// Returns `DagError::ResultNotFound` if:
/// - The task hasn't been executed yet
/// - The handle is invalid
///
/// # Examples
///
/// ```no_run
/// # use dagx::{task, DagRunner, Task};
/// #
/// struct Configuration {
/// setting: i32,
/// }
///
/// impl Configuration {
/// fn new(setting: i32) -> Self {
/// Self { setting }
/// }
/// }
///
/// #[task]
/// impl Configuration {
/// async fn run(&mut self) -> i32 { self.setting }
/// }
///
/// # async {
/// let mut dag = DagRunner::new();
///
/// // Construct task with specific setting value
/// let task = dag.add_task(Configuration::new(42));
///
///let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
///
/// assert_eq!(output.get(task).unwrap(), 42);
/// # };
/// ```
pub fn get<T: 'static + Send + Sync, H>(&mut self, handle: H) -> DagResult<T>
where
H: Into<TaskHandle<T>>,
{
let handle: TaskHandle<T> = handle.into();
let arc_output = self
.outputs
.remove(&handle.id)
.ok_or(DagError::ResultNotFound {
task_id: handle.id.0,
})?;
// Downcast Arc<dyn Any> to Arc<T>
let output = arc_output
.downcast::<T>()
.map_err(|_| DagError::TypeMismatch {
expected: std::any::type_name::<T>(),
found: "unknown",
})?;
Ok(Arc::into_inner(output).unwrap())
}
pub(super) fn new(
outputs: PassThroughHashMap<NodeId, Arc<dyn Any + Send + Sync + 'static>>,
) -> Self {
Self { outputs }
}
}
src/runner.rs
+40 -181
View File
@@ -10,34 +10,18 @@ use std::collections::{HashMap, HashSet, VecDeque};
use std::future::Future;
use std::hash::{BuildHasher, Hasher};
use std::panic::AssertUnwindSafe;
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};
use std::sync::Arc;
use futures_util::future::BoxFuture;
use futures_util::{stream::FuturesUnordered, FutureExt, StreamExt, TryFutureExt};
use parking_lot::Mutex;
#[cfg(feature = "tracing")]
use tracing::{debug, error, info, trace};
use crate::builder::TaskBuilder;
use crate::builder::{NodeId, TaskWire};
use crate::error::{DagError, DagResult};
use crate::node::{ExecutableNode, TypedNode};
use crate::task::Task;
use crate::types::{NodeId, Pending, TaskHandle};
// Guard to ensure run_lock is released even on early return or panic
struct RunGuard<'a> {
lock: &'a AtomicBool,
}
impl<'a> Drop for RunGuard<'a> {
fn drop(&mut self) {
self.lock.store(false, Ordering::SeqCst);
}
}
use crate::DagOutput;
#[derive(Default, Clone)]
pub(crate) struct PassThroughHasher {
@@ -76,10 +60,10 @@ pub(crate) type PassThroughHashMap<K, V> = HashMap<K, V, PassThroughHasher>;
/// # Workflow
///
/// 1. Create a new DAG with [`DagRunner::new`]
/// 2. Add tasks with [`DagRunner::add_task`] to get [`TaskBuilder`] builders
/// 3. Wire dependencies with [`TaskBuilder::depends_on`]
/// 2. Add tasks with [`DagRunner::add_task`] to get builders
/// 3. Wire dependencies with [`crate::TaskBuilder::depends_on`]
/// 4. Execute all tasks with [`DagRunner::run`]
/// 5. Optionally retrieve outputs with [`DagRunner::get`]
/// 5. Optionally retrieve outputs with [`DagOutput::get`]
///
/// # Examples
///
@@ -107,16 +91,16 @@ pub(crate) type PassThroughHashMap<K, V> = HashMap<K, V, PassThroughHasher>;
/// }
///
/// # async {
/// let dag = DagRunner::new();
/// let mut dag = DagRunner::new();
///
/// // Construct instances using ::new() pattern
/// let x = dag.add_task(LoadValue::new(2));
/// let y = dag.add_task(LoadValue::new(3));
/// let sum = dag.add_task(Add).depends_on((x, y));
///
/// dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
///let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
///
/// assert_eq!(dag.get(sum).unwrap(), 5);
/// assert_eq!(output.get(sum).unwrap(), 5);
/// # };
/// ```
///
@@ -127,13 +111,9 @@ pub(crate) type PassThroughHashMap<K, V> = HashMap<K, V, PassThroughHasher>;
/// Outputs are Arc-wrapped and stored separately for retrieval via get().
/// Arc enables efficient sharing during fanout without cloning data.
pub struct DagRunner {
pub(crate) nodes: Mutex<Vec<Option<Box<dyn ExecutableNode + Sync>>>>,
pub(crate) outputs:
Mutex<PassThroughHashMap<NodeId, std::sync::Arc<dyn std::any::Any + Send + Sync>>>,
pub(crate) edges: Mutex<PassThroughHashMap<NodeId, Vec<NodeId>>>, // node -> dependencies
pub(crate) dependents: Mutex<PassThroughHashMap<NodeId, Vec<NodeId>>>, // node -> tasks that depend on it
pub(crate) next_id: Mutex<u32>,
pub(crate) run_lock: AtomicBool, // Ensures only one run() at a time
pub(crate) nodes: Vec<Option<Box<dyn ExecutableNode + Sync>>>,
pub(crate) edges: PassThroughHashMap<NodeId, Vec<NodeId>>, // node -> dependencies
pub(crate) dependents: PassThroughHashMap<NodeId, Vec<NodeId>>, // node -> tasks that depend on it
}
impl Default for DagRunner {
@@ -150,31 +130,21 @@ impl DagRunner {
/// ```
/// use dagx::DagRunner;
///
/// let dag = DagRunner::new();
/// let mut dag = DagRunner::new();
/// ```
pub fn new() -> Self {
Self {
nodes: Mutex::new(Vec::new()),
outputs: Mutex::new(HashMap::default()),
edges: Mutex::new(HashMap::default()),
dependents: Mutex::new(HashMap::default()),
next_id: Mutex::new(0),
run_lock: AtomicBool::new(false),
nodes: Vec::new(),
edges: HashMap::default(),
dependents: HashMap::default(),
}
}
pub(crate) fn alloc_id(&self) -> NodeId {
let mut next_id = self.next_id.lock();
let id = NodeId(*next_id);
*next_id += 1;
id
}
/// Add a task instance to the DAG, returning a node builder for wiring dependencies.
///
/// The returned [`TaskBuilder<Tk, Pending>`](TaskBuilder) can be used to:
/// - Specify dependencies via [`TaskBuilder::depends_on`]
/// - Used directly as a [`TaskHandle`] to the task's output
/// If the task has no dependencies, a [`crate::TaskHandle`] will be returned.
/// If not, the returned [`crate::TaskBuilder`] can be used to
/// specify dependencies via [`crate::TaskBuilder::depends_on`].
///
/// # Examples
///
@@ -214,7 +184,7 @@ impl DagRunner {
/// }
///
/// # async {
/// let dag = DagRunner::new();
/// let mut dag = DagRunner::new();
///
/// // Construct task with initial value of 10
/// let base = dag.add_task(LoadValue::new(10));
@@ -222,17 +192,16 @@ impl DagRunner {
/// // Construct task with offset of 1
/// let inc = dag.add_task(AddOffset::new(1)).depends_on(base);
///
/// dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
/// assert_eq!(dag.get(inc).unwrap(), 11);
///let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
/// assert_eq!(output.get(inc).unwrap(), 11);
/// # };
/// ```
pub fn add_task<Tk>(&self, task: Tk) -> TaskBuilder<'_, Tk, Pending>
pub fn add_task<'dag, Input, Tk>(&'dag mut self, task: Tk) -> Tk::Retval<'dag>
where
Tk: Task + Sync + 'static,
Tk::Input: 'static,
Tk::Output: 'static,
Tk: TaskWire<Input>,
Input: Send + Sync + 'static,
{
let id = self.alloc_id();
let id = NodeId(self.nodes.len() as u32);
#[cfg(feature = "tracing")]
debug!(
@@ -242,15 +211,11 @@ impl DagRunner {
);
let node = TypedNode::new(task);
self.nodes.lock().push(Some(Box::new(node)));
self.edges.lock().insert(id, Vec::new());
self.dependents.lock().insert(id, Vec::new());
self.nodes.push(Some(Box::new(node)));
self.edges.insert(id, Vec::new());
self.dependents.insert(id, Vec::new());
TaskBuilder {
id,
dag: self,
_phantom: std::marker::PhantomData,
}
Tk::new_from_dag(id, self)
}
/// Run the entire DAG to completion using the provided spawner.
@@ -299,41 +264,24 @@ impl DagRunner {
/// }
///
/// # async {
/// let dag = DagRunner::new();
/// let mut dag = DagRunner::new();
///
/// let a = dag.add_task(Value(1));
/// let b = dag.add_task(Value(2));
/// let sum = dag.add_task(Add).depends_on((a, b));
///
/// dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap(); // Executes all tasks
///let mut output = dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap(); // Executes all tasks
/// # };
/// ```
#[inline]
#[cfg_attr(feature = "tracing", tracing::instrument(skip(self, spawner)))]
pub async fn run<S, F>(&self, spawner: S) -> DagResult<()>
pub async fn run<S, F>(mut self, spawner: S) -> DagResult<DagOutput>
where
S: Fn(BoxFuture<'static, DagResult<Arc<dyn Any + Send + Sync>>>) -> F,
F: Future<Output = DagResult<Arc<dyn Any + Send + Sync>>>,
{
#[cfg(feature = "tracing")]
info!("starting DAG execution");
// Acquire run lock to prevent concurrent executions
// Use atomic compare_exchange to check if already running
if self
.run_lock
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_err()
{
#[cfg(feature = "tracing")]
error!("DAG is already running - concurrent execution not supported");
return Err(DagError::ConcurrentExecution);
}
// Guard ensures lock is released on drop (even on early return or panic)
let _run_guard = RunGuard {
lock: &self.run_lock,
};
// Build topological layers
let layers = self.compute_layers()?;
@@ -346,8 +294,6 @@ impl DagRunner {
total_tasks, "computed topological layers"
);
let edges = self.edges.lock().clone();
let mut outputs: PassThroughHashMap<NodeId, Arc<dyn Any + Send + Sync>> =
HashMap::with_capacity_and_hasher(total_tasks, PassThroughHasher::default());
let mut first_error = None;
@@ -392,13 +338,10 @@ impl DagRunner {
);
// Take ownership of the node
let node = {
let mut nodes_lock = self.nodes.lock();
nodes_lock[node_id.0 as usize].take()
};
let node = self.nodes[node_id.0 as usize].take();
if let Some(node) = node {
let dependencies: Vec<_> = edges[&node_id]
let dependencies: Vec<_> = self.edges[&node_id]
.iter()
.flat_map(|dep| outputs.get(dep))
.cloned()
@@ -456,12 +399,9 @@ impl DagRunner {
trace!(task_id = node_id.0, "spawning task");
// Take ownership of the node
let node = {
let mut nodes_lock = self.nodes.lock();
nodes_lock[node_id.0 as usize].take()
};
let node = self.nodes[node_id.0 as usize].take();
if let Some(node) = node {
let dependencies: Vec<_> = edges[&node_id]
let dependencies: Vec<_> = self.edges[&node_id]
.iter()
.flat_map(|dep| outputs.get(dep))
.cloned()
@@ -522,13 +462,6 @@ impl DagRunner {
}
}
// Insert all successful outputs (avoiding holding lock across await)
let mut outputs_lock = self.outputs.lock();
for (node_id, output) in outputs {
outputs_lock.insert(node_id, output);
}
drop(outputs_lock);
// Return first error if any
if let Some(err) = first_error {
#[cfg(feature = "tracing")]
@@ -539,73 +472,7 @@ impl DagRunner {
#[cfg(feature = "tracing")]
info!("DAG execution completed successfully");
Ok(())
}
/// Retrieve a task's output after [`DagRunner::run`].
///
/// # Behavior
///
/// All task outputs are stored after execution and can be retrieved via get().
/// Each task's output can only be taken once, after which it will return `DagError::ResultNotFound`.
///
/// # Errors
///
/// Returns `DagError::ResultNotFound` if:
/// - The task hasn't been executed yet
/// - The handle is invalid
///
/// # Examples
///
/// ```no_run
/// # use dagx::{task, DagRunner, Task};
/// #
/// struct Configuration {
/// setting: i32,
/// }
///
/// impl Configuration {
/// fn new(setting: i32) -> Self {
/// Self { setting }
/// }
/// }
///
/// #[task]
/// impl Configuration {
/// async fn run(&mut self) -> i32 { self.setting }
/// }
///
/// # async {
/// let dag = DagRunner::new();
///
/// // Construct task with specific setting value
/// let task = dag.add_task(Configuration::new(42));
///
/// dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
///
/// assert_eq!(dag.get(task).unwrap(), 42);
/// # };
/// ```
pub fn get<T: 'static + Send + Sync, H>(&self, handle: H) -> DagResult<T>
where
H: Into<TaskHandle<T>>,
{
let handle: TaskHandle<T> = handle.into();
let mut outputs = self.outputs.lock();
let arc_output = outputs.remove(&handle.id).ok_or(DagError::ResultNotFound {
task_id: handle.id.0,
})?;
// Downcast Arc<dyn Any> to Arc<T>
let output = arc_output
.downcast::<T>()
.map_err(|_| DagError::TypeMismatch {
expected: std::any::type_name::<T>(),
found: "unknown",
})?;
Ok(Arc::into_inner(output).unwrap())
Ok(DagOutput::new(outputs))
}
#[inline]
@@ -617,18 +484,11 @@ impl DagRunner {
let mut layers = Vec::new();
// Calculate in-degrees: for each node, count how many dependencies it has
let edges = self.edges.lock();
// Compiler optimization hint: provides bounds information for HashSet sizing.
// This variable guides the optimizer's understanding of graph size, improving
// hash table and iterator performance in the topological sort below.
let _total_nodes = edges.len();
for (&node, deps) in edges.iter() {
for (&node, deps) in self.edges.iter() {
let degree = deps.len();
in_degree.insert(node, degree);
}
drop(edges); // Release lock early
// Find all nodes with in-degree 0 (sources - nodes with no dependencies)
let mut queue: VecDeque<NodeId> = in_degree
@@ -653,8 +513,7 @@ impl DagRunner {
visited.insert(node);
// For each node that depends on the current node
let dependents = self.dependents.lock();
if let Some(deps) = dependents.get(&node) {
if let Some(deps) = self.dependents.get(&node) {
for &dependent in deps {
if let Some(degree) = in_degree.get_mut(&dependent) {
*degree -= 1;
src/runner/tests.rs
+23 -80
View File
@@ -3,9 +3,10 @@
use std::collections::HashMap;
use std::hash::Hasher;
use crate::builder::TaskHandle;
use crate::error::DagError;
use crate::runner::{DagRunner, PassThroughHasher};
use crate::types::TaskHandle;
use crate::DagOutput;
// Initialize tracing subscriber for tests (idempotent)
#[cfg(feature = "tracing")]
@@ -40,16 +41,16 @@ impl TestTask {
#[test]
fn test_dag_runner_new() {
// Test that DagRunner::new() creates a new instance
let dag = DagRunner::new();
let mut dag = DagRunner::new();
// Initial state should be empty
assert_eq!(dag.nodes.lock().len(), 0);
assert_eq!(dag.edges.lock().len(), 0);
assert_eq!(dag.dependents.lock().len(), 0);
assert_eq!(dag.nodes.len(), 0);
assert_eq!(dag.edges.len(), 0);
assert_eq!(dag.dependents.len(), 0);
// Adding a task works
dag.add_task(TestTask { value: 42 });
assert_eq!(dag.nodes.lock().len(), 1);
assert_eq!(dag.nodes.len(), 1);
}
#[test]
@@ -58,19 +59,20 @@ fn test_dag_runner_default() {
let dag = DagRunner::default();
// Should behave the same as new()
assert_eq!(dag.nodes.lock().len(), 0);
assert_eq!(dag.edges.lock().len(), 0);
assert_eq!(dag.dependents.lock().len(), 0);
assert_eq!(dag.nodes.len(), 0);
assert_eq!(dag.edges.len(), 0);
assert_eq!(dag.dependents.len(), 0);
}
#[tokio::test]
async fn test_get_wrong_type() {
init_tracing();
// Test getting with wrong type - this should return TypeMismatch
let dag = DagRunner::new();
let mut dag = DagRunner::new();
let handle = dag.add_task(TestTask { value: 42 });
dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
let mut output = dag
.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
.await
.unwrap();
@@ -80,7 +82,7 @@ async fn test_get_wrong_type() {
_phantom: std::marker::PhantomData,
};
let result = dag.get(fake_handle);
let result = output.get(fake_handle);
assert!(result.is_err());
// When downcast fails, we get TypeMismatch
@@ -96,11 +98,12 @@ async fn test_get_wrong_type() {
async fn test_get_result_not_found() {
init_tracing();
// Test getting a result before running the DAG
let dag = DagRunner::new();
let mut dag = DagRunner::new();
let handle = dag.add_task(TestTask { value: 42 });
let handle_id = handle.id.0; // Save ID before moving handle
let mut output = DagOutput::new(HashMap::default());
let result = dag.get(handle);
let result = output.get(handle);
assert!(result.is_err());
match result.unwrap_err() {
DagError::ResultNotFound { task_id } => {
@@ -110,67 +113,6 @@ async fn test_get_result_not_found() {
}
}
#[tokio::test]
async fn test_concurrent_run_protection() {
init_tracing();
// Test that run_lock prevents concurrent runs
use std::sync::Arc;
use std::time::Duration;
use tokio::time::sleep;
let dag = DagRunner::new();
// Add a task that takes some time
struct SlowTask;
#[crate::task]
impl SlowTask {
async fn run(&self) -> i32 {
sleep(Duration::from_millis(100)).await;
42
}
}
dag.add_task(SlowTask);
// Wrap in Arc for sharing between tasks
let dag = Arc::new(dag);
let dag1 = Arc::clone(&dag);
let dag2 = Arc::clone(&dag);
// Use a barrier to ensure both runs start at exactly the same time
let barrier = Arc::new(tokio::sync::Barrier::new(2));
let barrier1 = barrier.clone();
let barrier2 = barrier.clone();
// Start two runs concurrently
let handle1 = tokio::spawn(async move {
barrier1.wait().await;
dag1.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
.await
});
let handle2 = tokio::spawn(async move {
barrier2.wait().await;
dag2.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
.await
});
// One should succeed, one should fail (concurrent execution not supported)
let result1 = handle1.await.unwrap();
let result2 = handle2.await.unwrap();
// Exactly one should be ok, one should be err
assert!(result1.is_ok() != result2.is_ok());
// The error should be about concurrent execution
if let Err(e) = result1 {
assert!(matches!(e, crate::error::DagError::ConcurrentExecution));
}
if let Err(e) = result2 {
assert!(matches!(e, crate::error::DagError::ConcurrentExecution));
}
}
#[tokio::test]
async fn test_invalid_node_id_in_get() {
init_tracing();
@@ -179,15 +121,16 @@ async fn test_invalid_node_id_in_get() {
// Create a handle with an ID that doesn't exist
let invalid_handle: TaskHandle<i32> = TaskHandle {
id: crate::types::NodeId(999),
id: crate::builder::NodeId(999),
_phantom: std::marker::PhantomData,
};
dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
let mut output = dag
.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
.await
.unwrap();
let result = dag.get(invalid_handle);
let result = output.get(invalid_handle);
assert!(result.is_err());
// With centralized output storage, we get ResultNotFound for invalid node IDs
@@ -228,8 +171,8 @@ async fn test_task_panic_in_multi_task_layer() {
}
}
let dag = DagRunner::new();
let source: TaskHandle<_> = dag.add_task(Source).into();
let mut dag = DagRunner::new();
let source = dag.add_task(Source);
// Create two tasks in the same layer - one panics, one doesn't
let _panic_task = dag.add_task(PanicTask).depends_on(source);
src/task.rs
+6 -4
View File
@@ -92,11 +92,13 @@ use std::{any::Any, future::Future, marker::PhantomData, slice::Iter, sync::Arc}
///
/// Downstream tasks receive `&T` after ExtractInput extracts from the Arc. This is transparent
/// to your task code - you work with `T`, the framework handles Arc wrapping/unwrapping.
pub trait Task: Send {
type Input: Send;
type Output: Send + Sync; // Sync required for Arc-wrapping and cross-thread sharing
pub trait Task<Input>: Send
where
Input: Send + Sync + 'static,
{
type Output: Send + Sync + 'static; // Sync required for Arc-wrapping and cross-thread sharing
fn run(self, input: TaskInput<Self::Input>) -> impl Future<Output = Self::Output> + Send;
fn run(self, input: TaskInput<Input>) -> impl Future<Output = Self::Output> + Send;
}
pub struct TaskInput<'inputs, Input: Send> {
src/task/tests.rs
+2 -3
View File
@@ -9,11 +9,10 @@ struct CustomTask {
state: i32,
}
impl Task for CustomTask {
type Input = (i32,);
impl Task<(i32,)> for CustomTask {
type Output = i32;
async fn run(mut self, input: TaskInput<'_, Self::Input>) -> i32 {
async fn run(mut self, input: TaskInput<'_, (i32,)>) -> i32 {
self.state += input.next().0;
self.state
}
src/types.rs
-86
View File
@@ -1,86 +0,0 @@
//! Core type definitions for DAG nodes and handles.
//!
//! This module defines the fundamental types used throughout the DAG system,
//! including node identifiers, task handles, and type-state markers.
use crate::builder::TaskBuilder;
use crate::task::Task;
/// Opaque node identifier
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub struct NodeId(pub u32);
/// Opaque, typed token for a node's output.
///
/// A `TaskHandle<T>` provides compile-time type-safe access to a task's output.
/// You can:
/// 1. Pass it to [`crate::TaskBuilder::depends_on`] to wire up dependencies
/// 2. Use it with [`crate::DagRunner::get`] to retrieve the output after [`crate::DagRunner::run`]
///
/// Handles are cheap to clone and copy.
///
/// # Examples
///
/// ```no_run
/// # use dagx::{task, DagRunner, Task};
/// #
/// # struct LoadValue { value: i32 }
/// # impl LoadValue { pub fn new(v: i32) -> Self { Self { value: v } } }
/// # #[task]
/// # impl LoadValue {
/// # async fn run(&mut self) -> i32 { self.value }
/// # }
/// # async {
/// let dag = DagRunner::new();
/// let node = dag.add_task(LoadValue::new(42));
///
/// dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() }).await.unwrap();
///
/// assert_eq!(dag.get(node).unwrap(), 42);
/// # };
/// ```
pub struct TaskHandle<T> {
pub(crate) id: NodeId,
pub(crate) _phantom: std::marker::PhantomData<fn() -> T>,
}
impl<T> Clone for TaskHandle<T> {
fn clone(&self) -> Self {
*self
}
}
impl<T> Copy for TaskHandle<T> {}
/// Type-level marker for empty dependency list.
///
/// This type appears in the type signature of [`crate::TaskBuilder<Tk, Pending>`](crate::TaskBuilder) before
/// dependencies are specified with [`crate::TaskBuilder::depends_on`].
///
/// You don't interact with this type directly; it's part of the type-state
/// pattern that ensures dependencies are wired correctly at compile time.
pub struct Pending;
// TaskBuilder can be converted to a TaskHandle ONLY if Input is () (unit type)
// This enforces at compile time that tasks with non-unit inputs MUST call .depends_on()
impl<'a, Tk: Task<Input = ()>, Deps> From<TaskBuilder<'a, Tk, Deps>> for TaskHandle<Tk::Output> {
fn from(node: TaskBuilder<'a, Tk, Deps>) -> Self {
TaskHandle {
id: node.id,
_phantom: std::marker::PhantomData,
}
}
}
// TaskHandle can be converted from &TaskHandle (for .get() calls)
impl<T> From<&TaskHandle<T>> for TaskHandle<T> {
fn from(handle: &TaskHandle<T>) -> Self {
*handle
}
}
#[cfg(test)]
mod tests;
#[cfg(test)]
mod coverage_tests;
src/types/tests.rs
-50
View File
@@ -1,50 +0,0 @@
//! Unit tests for types module
use crate::types::{NodeId, TaskHandle};
use std::marker::PhantomData;
#[test]
#[allow(clippy::clone_on_copy)]
fn test_task_handle_clone() {
let handle: TaskHandle<(i32,)> = TaskHandle {
id: NodeId(42),
_phantom: PhantomData,
};
// Test Clone implementation
let cloned = handle.clone();
assert_eq!(handle.id.0, cloned.id.0);
let handle2 = handle.clone();
let handle3 = handle2.clone();
assert_eq!(handle3.id.0, 42);
}
#[test]
fn test_task_handle_copy() {
let handle: TaskHandle<String> = TaskHandle {
id: NodeId(99),
_phantom: PhantomData,
};
// Test that TaskHandle implements Copy
let copied = handle;
let copied2 = handle; // This wouldn't compile if not Copy
assert_eq!(copied.id.0, 99);
assert_eq!(copied2.id.0, 99);
}
#[test]
fn test_task_handle_ref_from_conversion() {
// Test From<&TaskHandle<T>> for TaskHandle<T> (line 90-93)
let handle: TaskHandle<String> = TaskHandle {
id: NodeId(123),
_phantom: PhantomData,
};
let handle_ref = &handle;
let converted: TaskHandle<String> = handle_ref.into();
assert_eq!(converted.id.0, 123);
assert_eq!(handle.id, converted.id);
}