diff --git a/Cargo.lock b/Cargo.lock index ce288a7..e680cbf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -442,6 +442,7 @@ dependencies = [ "criterion", "dagrs", "dagx-macros", + "dashmap", "futures", "futures-executor", "parking_lot", @@ -461,6 +462,20 @@ dependencies = [ "syn 2.0.106", ] +[[package]] +name = "dashmap" +version = "6.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "either" version = "1.15.0" @@ -641,6 +656,12 @@ dependencies = [ "crunchy", ] +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + [[package]] name = "hermit-abi" version = "0.5.2" diff --git a/Cargo.toml b/Cargo.toml index 55c9914..ed84368 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,16 +22,25 @@ default = [] tracing = ["dep:tracing"] [lints.rust] -unexpected_cfgs = { level = "warn", check-cfg = ['cfg(tarpaulin)', 'cfg(tarpaulin_include)'] } +unexpected_cfgs = { level = "warn", check-cfg = [ + 'cfg(tarpaulin)', + 'cfg(tarpaulin_include)', +] } [dependencies] futures = "0.3" parking_lot = "0.12" dagx-macros = { version = "0.3.1", path = "./dagx-macros" } tracing = { version = "0.1", optional = true } +dashmap = "6.1.0" [dev-dependencies] -tokio = { version = "1", features = ["macros", "rt-multi-thread", "time", "sync"] } +tokio = { version = "1", features = [ + "macros", + "rt-multi-thread", + "time", + "sync", +] } async-std = { version = "1", features = ["attributes"] } smol = "2" criterion = "0.7" diff --git a/README.md b/README.md index 4cc971f..2272c57 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ See [how it works](docs/CYCLE_PREVENTION.md). ```rust let sum = dag.add_task(Add).depends_on((&x, &y)); -dag.run(|fut| tokio::spawn(fut)).await?; +dag.run().await?; ``` That's it. No trait boilerplate, no manual channels, no node IDs. @@ -90,7 +90,7 @@ async fn main() { let sum = dag.add_task(Add).depends_on((&x, &y)); // Execute with true parallelism - dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); + dag.run().await.unwrap(); // Retrieve results assert_eq!(dag.get(sum).unwrap(), 5); @@ -223,13 +223,13 @@ dagx works with any async runtime. Provide a spawner function to `run()`: ```rust // With Tokio -dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); +dag.run().await.unwrap(); // With async-std -dag.run(|fut| { async_std::task::spawn(fut); }).await.unwrap(); +dag.run().await.unwrap(); // With smol -dag.run(|fut| { smol::spawn(fut).detach(); }).await.unwrap(); +dag.run().await.unwrap(); ``` ## Tutorials & Examples diff --git a/dagx-macros/src/lib.rs b/dagx-macros/src/lib.rs index f7d7fb5..4339bb3 100644 --- a/dagx-macros/src/lib.rs +++ b/dagx-macros/src/lib.rs @@ -335,7 +335,7 @@ pub fn task(_attr: TokenStream, item: TokenStream) -> TokenStream { quote! { fn extract_and_run( self, - _receivers: Vec>, + _dependencies: Vec>, ) -> impl std::future::Future> + Send { async move { let input = (); @@ -350,24 +350,20 @@ pub fn task(_attr: TokenStream, item: TokenStream) -> TokenStream { quote! { fn extract_and_run( self, - mut receivers: Vec>, + mut dependencies: Vec>, ) -> impl std::future::Future> + Send { async move { - use futures::channel::oneshot; use std::sync::Arc; - if receivers.len() != 1 { - return Err(format!("Expected 1 dependency, got {}", receivers.len())); + if dependencies.len() != 1 { + return Err(format!("Expected 1 dependency, got {}", dependencies.len())); } - let rx = *receivers.pop() + let arc_value = dependencies.pop() .unwrap() - .downcast::>>() + .downcast::<#param_type>() .map_err(|_| format!("Type mismatch: expected Arc<{}>", std::any::type_name::<#param_type>()))?; - let arc_value = rx.await - .map_err(|_| "Channel closed before receiving value".to_string())?; - let input = (*arc_value).clone(); Ok(self.run(input).await) } @@ -380,51 +376,36 @@ pub fn task(_attr: TokenStream, item: TokenStream) -> TokenStream { let param_types: Vec<_> = params.iter().map(|(_, ty)| ty).collect(); let indices: Vec<_> = (0..param_count).collect(); - // Generate unique receiver variable names - let rx_vars: Vec<_> = (0..param_count) - .map(|i| syn::Ident::new(&format!("rx_{}", i), proc_macro2::Span::call_site())) - .collect(); - // Create syn::Index for tuple field access (avoids the suffix warning) let syn_indices: Vec<_> = (0..param_count).map(syn::Index::from).collect(); quote! { fn extract_and_run( self, - receivers: Vec>, + dependencies: Vec>, ) -> impl std::future::Future> + Send { async move { - use futures::channel::oneshot; use std::sync::Arc; let expected_count = #param_count; - if receivers.len() != expected_count { - return Err(format!("Expected {} dependencies, got {}", expected_count, receivers.len())); + if dependencies.len() != expected_count { + return Err(format!("Expected {} dependencies, got {}", expected_count, dependencies.len())); } - let mut iter = receivers.into_iter(); + let mut iter = dependencies.into_iter(); // Extract each receiver - #( - let #rx_vars = *iter.next() + let arc_results = (#( + iter.next() .ok_or_else(|| format!("Missing receiver at index {}", #indices))? - .downcast::>>() + .downcast::<#param_types>() .map_err(|_| format!("Type mismatch at index {}: expected Arc<{}>", - #indices, std::any::type_name::<#param_types>()))?; - )* - - // Await all channels concurrently - let arc_results = futures::join!( - #( - async move { - #rx_vars.await.map_err(|_| format!("Channel {} closed", #indices)) - } - ),* - ); + #indices, std::any::type_name::<#param_types>()))?, + )*); // Clone inner values and build tuple let input = (#( - (*arc_results.#syn_indices?).clone() + (*arc_results.#syn_indices).clone() ),*); Ok(self.run(input).await) diff --git a/docs/COMPARISONS.md b/docs/COMPARISONS.md index ad30468..2c53ccc 100644 --- a/docs/COMPARISONS.md +++ b/docs/COMPARISONS.md @@ -19,7 +19,7 @@ let 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)); -dag.run(|fut| tokio::spawn(fut)).await?; +dag.run().await?; ``` **dagrs** (Complex: traits + channels + IDs + manual wiring): diff --git a/docs/CYCLE_PREVENTION.md b/docs/CYCLE_PREVENTION.md index fbba17e..839609c 100644 --- a/docs/CYCLE_PREVENTION.md +++ b/docs/CYCLE_PREVENTION.md @@ -13,7 +13,7 @@ let a = dag.add_task(TaskA); let b = dag.add_task(TaskB); a.depends_on(b); b.depends_on(a); // Compiles fine... -dag.run().await?; // ❌ Error: "Cycle detected!" (at runtime) +dag.run() ``` **dagx prevents this at compile time:** diff --git a/src/builder.rs b/src/builder.rs index d5bba4e..b4e6b40 100644 --- a/src/builder.rs +++ b/src/builder.rs @@ -65,7 +65,7 @@ impl IsUnitType for () {} /// let b = b.depends_on(&a); /// // Now b is a TaskHandle /// -/// dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); +/// dag.run().await.unwrap(); /// assert_eq!(dag.get(&b).unwrap(), 20); /// # }; /// ``` @@ -125,7 +125,7 @@ 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| { tokio::spawn(fut); }).await.unwrap(); + /// dag.run().await.unwrap(); /// # }; /// ``` pub fn depends_on(self, deps: D) -> TaskHandle diff --git a/src/error.rs b/src/error.rs index b7afd57..7cbb01a 100644 --- a/src/error.rs +++ b/src/error.rs @@ -63,7 +63,7 @@ impl std::fmt::Display for DagError { f, "Result not found for task #{}.\n\ \n\ - Call dag.run() before accessing results with get().", + Call dag.run().", task_id ) } diff --git a/src/extract.rs b/src/extract.rs index f1b1b54..7c9f06f 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -11,7 +11,7 @@ //! //! # How it works //! -//! When a task executes, it needs to extract its inputs from type-erased oneshot receivers. +//! When a task executes, it needs to extract its inputs from type-erased oneshot dependencies. //! Outputs are wrapped in Arc for efficient fanout (cheap Arc clones instead of data clones). //! //! **Modern approach (via `#[task]` macro)**: @@ -35,8 +35,6 @@ use std::any::Any; use std::sync::Arc; -use futures::channel::oneshot; - /// Helper trait for async input extraction from oneshot channels. /// /// Type erasure occurs only at the ExecutableNode trait boundary - by the time @@ -47,14 +45,14 @@ use futures::channel::oneshot; /// which allows ANY type to work without requiring ExtractInput implementations. #[allow(dead_code)] pub(crate) trait ExtractInput: Sized + Clone { - fn extract_from_channels( - receivers: Vec>, + fn extract_from_deps( + deps: Vec>, ) -> impl std::future::Future> + Send; } // Unit type - no dependencies impl ExtractInput for () { - async fn extract_from_channels(_receivers: Vec>) -> Result { + async fn extract_from_deps(_deps: Vec>) -> Result { Ok(()) } } @@ -67,15 +65,14 @@ macro_rules! impl_extract_single { ($($t:ty),+) => { $( impl ExtractInput for $t { - async fn extract_from_channels(mut receivers: Vec>) -> Result { - if receivers.len() != 1 { - return Err(format!("Expected 1 dependency, got {}", receivers.len())); + async fn extract_from_deps(mut dependencies: Vec>) -> Result { + if dependencies.len() != 1 { + return Err(format!("Expected 1 dependency, got {}", dependencies.len())); } - let rx = *receivers.pop() + let arc_value = dependencies.pop() .unwrap() - .downcast::>>() + .downcast::<$t>() .map_err(|_| format!("Type mismatch: expected Arc<{}>", std::any::type_name::<$t>()))?; - let arc_value = rx.await.map_err(|_| "Channel closed before receiving value".to_string())?; Ok((*arc_value).clone()) } } @@ -93,16 +90,16 @@ where T: 'static + Clone + Send + Sync, E: 'static + Clone + Send + Sync, { - async fn extract_from_channels( - mut receivers: Vec>, + async fn extract_from_deps( + mut dependencies: Vec>, ) -> Result { - if receivers.len() != 1 { - return Err(format!("Expected 1 dependency, got {}", receivers.len())); + if dependencies.len() != 1 { + return Err(format!("Expected 1 dependency, got {}", dependencies.len())); } - let rx = *receivers + let arc_value = dependencies .pop() .unwrap() - .downcast::>>>() + .downcast::>() .map_err(|_| { format!( "Type mismatch: expected Arc>", @@ -110,9 +107,6 @@ where std::any::type_name::() ) })?; - let arc_value = rx - .await - .map_err(|_| "Channel closed before receiving value".to_string())?; Ok((*arc_value).clone()) } } @@ -122,25 +116,22 @@ impl ExtractInput for Option where T: 'static + Clone + Send + Sync, { - async fn extract_from_channels( - mut receivers: Vec>, + async fn extract_from_deps( + mut dependencies: Vec>, ) -> Result { - if receivers.len() != 1 { - return Err(format!("Expected 1 dependency, got {}", receivers.len())); + if dependencies.len() != 1 { + return Err(format!("Expected 1 dependency, got {}", dependencies.len())); } - let rx = *receivers + let arc_value = dependencies .pop() .unwrap() - .downcast::>>>() + .downcast::>() .map_err(|_| { format!( "Type mismatch: expected Arc>", std::any::type_name::() ) })?; - let arc_value = rx - .await - .map_err(|_| "Channel closed before receiving value".to_string())?; Ok((*arc_value).clone()) } } @@ -150,25 +141,22 @@ impl ExtractInput for Vec where T: 'static + Clone + Send + Sync, { - async fn extract_from_channels( - mut receivers: Vec>, + async fn extract_from_deps( + mut dependencies: Vec>, ) -> Result { - if receivers.len() != 1 { - return Err(format!("Expected 1 dependency, got {}", receivers.len())); + if dependencies.len() != 1 { + return Err(format!("Expected 1 dependency, got {}", dependencies.len())); } - let rx = *receivers + let arc_value = dependencies .pop() .unwrap() - .downcast::>>>() + .downcast::>() .map_err(|_| { format!( "Type mismatch: expected Arc>", std::any::type_name::() ) })?; - let arc_value = rx - .await - .map_err(|_| "Channel closed before receiving value".to_string())?; Ok((*arc_value).clone()) } } @@ -179,16 +167,16 @@ where K: 'static + Clone + Send + Sync + Eq + std::hash::Hash, V: 'static + Clone + Send + Sync, { - async fn extract_from_channels( - mut receivers: Vec>, + async fn extract_from_deps( + mut dependencies: Vec>, ) -> Result { - if receivers.len() != 1 { - return Err(format!("Expected 1 dependency, got {}", receivers.len())); + if dependencies.len() != 1 { + return Err(format!("Expected 1 dependency, got {}", dependencies.len())); } - let rx = *receivers + let arc_value = dependencies .pop() .unwrap() - .downcast::>>>() + .downcast::>() .map_err(|_| { format!( "Type mismatch: expected Arc>", @@ -196,9 +184,6 @@ where std::any::type_name::() ) })?; - let arc_value = rx - .await - .map_err(|_| "Channel closed before receiving value".to_string())?; Ok((*arc_value).clone()) } } @@ -209,25 +194,22 @@ impl ExtractInput for Arc where T: 'static + Clone + Send + Sync, { - async fn extract_from_channels( - mut receivers: Vec>, + async fn extract_from_deps( + mut dependencies: Vec>, ) -> Result { - if receivers.len() != 1 { - return Err(format!("Expected 1 dependency, got {}", receivers.len())); + if dependencies.len() != 1 { + return Err(format!("Expected 1 dependency, got {}", dependencies.len())); } - let rx = *receivers + let arc_arc_value = dependencies .pop() .unwrap() - .downcast::>>>() + .downcast::>() .map_err(|_| { format!( "Type mismatch: expected Arc>", std::any::type_name::() ) })?; - let arc_arc_value = rx - .await - .map_err(|_| "Channel closed before receiving value".to_string())?; // Return the inner Arc (clone is cheap - just refcount increment) Ok(Arc::clone(&*arc_arc_value)) } @@ -251,43 +233,32 @@ where macro_rules! impl_extract_tuple { ($($T:ident : $idx:tt),+) => { impl<$($T: 'static + Clone + Send + Sync),+> ExtractInput for ($($T,)+) { - async fn extract_from_channels(receivers: Vec>) -> Result { - // Verify we have the correct number of receivers + async fn extract_from_deps(dependencies: Vec>) -> Result { + // Verify we have the correct number of dependencies let expected_count = [$($idx),+].len(); - if receivers.len() != expected_count { - return Err(format!("Expected {} dependencies, got {}", expected_count, receivers.len())); + if dependencies.len() != expected_count { + return Err(format!("Expected {} dependencies, got {}", expected_count, dependencies.len())); } // Convert Vec to iterator so we can consume it - let mut iter = receivers.into_iter(); + let mut iter = dependencies.into_iter(); - // Extract each receiver by consuming from the iterator - // Receivers now contain Arc-wrapped values for efficient fanout + // dependencies now contain Arc-wrapped values for efficient fanout // Dereference Box after downcasting to get the receiver itself // Use type names as variable names to avoid numeric suffix issues - $( - #[allow(non_snake_case)] - let $T = *iter.next() - .ok_or_else(|| format!("Missing receiver at index {}", $idx))? - .downcast::>>() - .map_err(|_| format!("Type mismatch at index {}: expected Arc<{}>", - $idx, std::any::type_name::<$T>()))?; - )+ - - // Await all channels concurrently - this is a key optimization! - // We can start receiving as soon as ANY dependency completes, not waiting - // for them sequentially. - let arc_results = futures::join!( + let arc_results = ( $( - async move { - $T.await.map_err(|_| format!("Channel {} closed", $idx)) - } - ),+ + iter.next() + .ok_or_else(|| format!("Missing receiver at index {}", $idx))? + .downcast::<$T>() + .map_err(|_| format!("Type mismatch at index {}: expected Arc<{}>", + $idx, std::any::type_name::<$T>()))?, + )+ ); // Clone inner values from Arc (cheap for small types, necessary for ownership) Ok(($( - (*arc_results.$idx?).clone() + (*arc_results.$idx).clone() ,)+)) } } diff --git a/src/lib.rs b/src/lib.rs index 5a4660d..903fe3b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -52,7 +52,7 @@ //! let sum = dag.add_task(Add).depends_on((&x, &y)); //! //! // Execute and retrieve results -//! dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); +//! dag.run().await.unwrap(); //! assert_eq!(dag.get(sum).unwrap(), 5); //! # }; //! ``` @@ -322,7 +322,7 @@ //! let plus1 = dag.add_task(Add(1)).depends_on(&base); //! let times2 = dag.add_task(Scale(2)).depends_on(&base); //! -//! dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); +//! dag.run().await.unwrap(); //! //! assert_eq!(dag.get(plus1).unwrap(), 11); //! assert_eq!(dag.get(times2).unwrap(), 20); @@ -378,7 +378,7 @@ //! let active = dag.add_task(Active(true)); //! let result = dag.add_task(FormatUser).depends_on((&name, &age, &active)); //! -//! dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); +//! dag.run().await.unwrap(); //! //! assert_eq!(dag.get(result).unwrap(), "User: Alice, Age: 30, Active: true"); //! # }; @@ -430,7 +430,7 @@ //! // Layer 3: Final result //! let total = dag.add_task(Add).depends_on((&sum_xy, &prod_yz)); // 5 + 15 = 20 //! -//! dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); +//! dag.run().await.unwrap(); //! //! assert_eq!(dag.get(total).unwrap(), 20); //! # }; @@ -473,7 +473,7 @@ //! let user = dag.add_task(CreateUser); //! let formatted = dag.add_task(FormatUser).depends_on(&user); //! -//! dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); +//! dag.run().await.unwrap(); //! //! assert_eq!(dag.get(formatted).unwrap(), "Alice is 30 years old"); //! # }; @@ -539,11 +539,11 @@ //! # let dag = DagRunner::new(); //! # let node = dag.add_task(Value(42)); //! // Simple approach with .unwrap() -//! dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); +//! dag.run().await.unwrap(); //! let result = dag.get(node).unwrap(); //! //! // Or handle errors explicitly -//! match dag.run(|fut| { tokio::spawn(fut); }).await { +//! match dag.run().await { //! Ok(_) => println!("DAG executed successfully"), //! Err(e) => eprintln!("DAG execution failed: {}", e), //! } @@ -656,7 +656,7 @@ //! let task2 = dag.add_task(ProcessData).depends_on(&data); //! let task3 = dag.add_task(ProcessData).depends_on(&data); //! -//! dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); +//! dag.run().await.unwrap(); //! # }; //! ``` //! diff --git a/src/node.rs b/src/node.rs index ee0a24c..41d3b27 100644 --- a/src/node.rs +++ b/src/node.rs @@ -11,18 +11,10 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use futures::channel::oneshot; - use crate::error::DagError; use crate::task::Task; use crate::types::NodeId; -/// Type-erased channel endpoints (senders or receivers) -type ErasedChannels = Vec>; - -/// Type-erased channel pair (senders, receivers) -type ChannelPair = (ErasedChannels, ErasedChannels); - /// Type-erased async execution result (Arc-wrapped for efficient fanout) type ExecuteFuture = Pin, DagError>> + Send>>; @@ -38,21 +30,13 @@ type ExecuteFuture = /// /// This is an implementation detail and not part of the public API. pub(crate) trait ExecutableNode: Send { - /// Create type-erased output channels for this node's dependents. - /// - /// Returns (senders, receivers) where: - /// - senders: Type-erased senders that this node will use to send outputs - /// - receivers: Type-erased receivers that dependents will use to receive inputs - fn create_output_channels(&self, num_dependents: usize) -> ChannelPair; - /// Execute the task with type-erased input receivers and output senders. /// /// The TypedNode implementation knows the concrete types and can safely downcast. /// Consumes the node and returns the output value. - fn execute_with_channels( + fn execute_with_deps( self: Box, - receivers: ErasedChannels, - senders: ErasedChannels, + dependencies: Vec>, ) -> ExecuteFuture; } @@ -86,30 +70,15 @@ where T::Input: 'static + Clone, T::Output: 'static + Clone, { - fn create_output_channels(&self, num_dependents: usize) -> ChannelPair { - let mut senders = Vec::new(); - let mut receivers = Vec::new(); - - for _ in 0..num_dependents { - // Create channels that pass Arc-wrapped outputs (efficient for fanout) - let (tx, rx) = oneshot::channel::>(); - senders.push(Box::new(tx) as Box); - receivers.push(Box::new(rx) as Box); - } - - (senders, receivers) - } - - fn execute_with_channels( + fn execute_with_deps( self: Box, - receivers: ErasedChannels, - senders: ErasedChannels, + dependencies: Vec>, ) -> ExecuteFuture { Box::pin(async move { // Use the task's extract_and_run method which has inline extraction logic // generated by the #[task] macro. This allows ANY type to work without // requiring ExtractInput trait implementations. - let output = match self.task.extract_and_run(receivers).await { + let output = match self.task.extract_and_run(dependencies).await { Ok(output) => output, Err(_msg) => { // Failed to extract input - likely channel closed or type mismatch @@ -124,17 +93,6 @@ where // Downstream tasks receive &T after extraction from Arc. let arc_output = Arc::new(output); - // Send Arc to all dependents via their channels - // TypedNode knows T::Output, so downcast is safe - for sender_any in senders { - if let Ok(sender) = sender_any.downcast::>>() { - // Clone Arc (just atomic refcount increment, NOT data clone) - // This is the key performance optimization for fan-out patterns - // Ignore send errors - dependent might have been cancelled - let _ = sender.send(Arc::clone(&arc_output)); - } - } - // Return Arc-wrapped output (for storage by runner) Ok(arc_output as Arc) }) diff --git a/src/runner.rs b/src/runner.rs index 65258bb..c1ab82c 100644 --- a/src/runner.rs +++ b/src/runner.rs @@ -5,6 +5,7 @@ //! //! Uses Mutex for interior mutability to enable builder pattern (`&self` instead of `&mut self`). +use std::any::Any; use std::collections::{HashMap, HashSet, VecDeque}; use std::panic::AssertUnwindSafe; use std::sync::{ @@ -12,9 +13,9 @@ use std::sync::{ Arc, }; -use futures::channel::mpsc; -use futures::future::BoxFuture; -use futures::{FutureExt, StreamExt}; +use dashmap::DashMap; +use futures::stream::FuturesUnordered; +use futures::{FutureExt, StreamExt, TryFutureExt}; use parking_lot::Mutex; #[cfg(feature = "tracing")] @@ -82,7 +83,7 @@ impl<'a> Drop for RunGuard<'a> { /// let y = dag.add_task(LoadValue::new(3)); /// let sum = dag.add_task(Add).depends_on((&x, &y)); /// -/// dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); +/// dag.run().await.unwrap(); /// /// assert_eq!(dag.get(sum).unwrap(), 5); /// # }; @@ -95,7 +96,7 @@ impl<'a> Drop for RunGuard<'a> { /// 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>>>, + pub(crate) nodes: Mutex>>>, pub(crate) outputs: Mutex>>, pub(crate) edges: Mutex>>, // node -> dependencies pub(crate) dependents: Mutex>>, // node -> tasks that depend on it @@ -188,13 +189,13 @@ impl DagRunner { /// // Construct task with offset of 1 /// let inc = dag.add_task(AddOffset::new(1)).depends_on(&base); /// - /// dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); + /// dag.run().await.unwrap(); /// assert_eq!(dag.get(&inc).unwrap(), 11); /// # }; /// ``` pub fn add_task(&self, task: Tk) -> TaskBuilder<'_, Tk, Pending> where - Tk: Task + 'static, + Tk: Task + Sync + 'static, Tk::Input: 'static + Clone, Tk::Output: 'static + Clone, { @@ -272,7 +273,7 @@ impl DagRunner { /// let b = dag.add_task(Value(2)); /// let sum = dag.add_task(Add).depends_on((&a, &b)); /// - /// dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); // Executes all tasks + /// dag.run().await.unwrap(); // Executes all tasks /// # }; /// ``` /// @@ -283,11 +284,8 @@ impl DagRunner { /// Type erasure occurs only at the ExecutableNode trait boundary - channels are created /// with full type information and type-erased just before passing to execute_with_channels. #[inline] - #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, spawner)))] - pub async fn run(&self, spawner: S) -> DagResult<()> - where - S: Fn(BoxFuture<'static, ()>), - { + #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))] + pub async fn run(&self) -> DagResult<()> { #[cfg(feature = "tracing")] info!("starting DAG execution"); // Acquire run lock to prevent concurrent executions @@ -320,47 +318,8 @@ impl DagRunner { let edges = self.edges.lock().clone(); - // Create all channels upfront - // Map: (producer_id, consumer_id) -> receiver index - // We store receivers in a Vec per consumer to maintain order - let mut consumer_receivers: HashMap>> = - HashMap::new(); - let mut producer_senders: HashMap>> = - HashMap::new(); - - // Create channels for each producer-consumer relationship - // We need to lock nodes temporarily to call create_output_channels - { - let nodes_lock = self.nodes.lock(); - - for (consumer_id, producer_ids) in &edges { - for &producer_id in producer_ids { - // Get the producer node (it's still in the vector at this point) - if let Some(Some(producer_node)) = nodes_lock.get(producer_id.0) { - // Ask the producer to create ONE channel for this consumer - let (mut senders, mut receivers) = producer_node.create_output_channels(1); - - // Store the sender for the producer - producer_senders - .entry(producer_id) - .or_default() - .push(senders.pop().unwrap()); - - // Store the receiver for the consumer (order matters!) - consumer_receivers - .entry(*consumer_id) - .or_default() - .push(receivers.pop().unwrap()); - } - } - } - } - - // Create a channel to collect Arc-wrapped outputs from all tasks - let (output_tx, mut output_rx) = mpsc::unbounded::<( - NodeId, - Result, DagError>, - )>(); + let outputs: DashMap> = DashMap::new(); + let mut first_error = None; for layer in layers { #[cfg(feature = "tracing")] @@ -402,8 +361,6 @@ impl DagRunner { "executing task inline (single-task layer optimization)" ); - let out_tx = output_tx.clone(); - // Take ownership of the node let node = { let mut nodes_lock = self.nodes.lock(); @@ -411,11 +368,11 @@ impl DagRunner { }; if let Some(node) = node { - // Take the receivers for this consumer - let receivers = consumer_receivers.remove(&node_id).unwrap_or_default(); - - // Take the senders for this producer - let senders = producer_senders.remove(&node_id).unwrap_or_default(); + let dependencies: Vec<_> = edges[&node_id] + .iter() + .flat_map(|dep| outputs.get(dep)) + .map(|value| value.clone()) + .collect(); // Execute inline with panic handling. // @@ -423,7 +380,7 @@ impl DagRunner { // DagError::TaskPanicked, matching the behavior of async runtimes when // they spawn tasks. This guarantees consistent error handling whether // a task executes inline (single-task layer) or spawned (multi-task layer). - let result = AssertUnwindSafe(node.execute_with_channels(receivers, senders)) + let result = AssertUnwindSafe(node.execute_with_deps(dependencies)) .catch_unwind() .await .unwrap_or_else(|panic_payload| { @@ -451,79 +408,94 @@ impl DagRunner { }); // Send the output (ignore send errors - receiver may be dropped) - let _ = out_tx.unbounded_send((node_id, result)); + match result { + Ok(output) => { + outputs.insert(node_id, output); + } + Err(e) => { + first_error.get_or_insert(e); + } + } } } else { // Slow path: Multiple tasks require spawning and coordination - // Create a channel to track task completion for this layer - let (tx, mut rx) = mpsc::channel::<()>(layer.len()); - // Spawn each task in this layer - for &node_id in &layer { - #[cfg(feature = "tracing")] - trace!(task_id = node_id.0, "spawning task"); - let mut task_tx = tx.clone(); - let out_tx = output_tx.clone(); + let outputs = &outputs; + let mut futures: FuturesUnordered<_> = FuturesUnordered::new(); + layer + .into_iter() + .flat_map(|node_id| { + #[cfg(feature = "tracing")] + 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].take() - }; - - if let Some(node) = node { - // Take the receivers for this consumer - let receivers = consumer_receivers.remove(&node_id).unwrap_or_default(); - - // Take the senders for this producer - let senders = producer_senders.remove(&node_id).unwrap_or_default(); - - // Create a 'static future by taking ownership - let inner_future = async move { - let task_future = node.execute_with_channels(receivers, senders); - - let result = task_future.await; - - // Send the output (ignore send errors - receiver may be dropped) - let _ = out_tx.unbounded_send((node_id, result)); - - // Signal completion (ignore send errors - receiver may be dropped) - let _ = task_tx.try_send(()); + // Take ownership of the node + let node = { + let mut nodes_lock = self.nodes.lock(); + nodes_lock[node_id.0].take() }; + if let Some(node) = node { + let dependencies: Vec<_> = edges[&node_id] + .iter() + .flat_map(|dep| outputs.get(dep)) + .map(|value| value.clone()) + .collect(); - // Spawn the task using the provided spawner - spawner(Box::pin(inner_future)); + // Spawn the task using the provided spawner + let inner_future = async move { + let result = node.execute_with_deps(dependencies).await?; + outputs.insert(node_id, result); + Ok(()) + }; + + Some( + AssertUnwindSafe(Box::pin(inner_future)) + .catch_unwind() + .unwrap_or_else(move |panic_payload| { + // Convert panic to error + let panic_message = + if let Some(s) = panic_payload.downcast_ref::<&str>() { + s.to_string() + } else if let Some(s) = + panic_payload.downcast_ref::() + { + s.clone() + } else { + "unknown panic".to_string() + }; + + #[cfg(feature = "tracing")] + error!( + task_id = node_id.0, + panic_message = %panic_message, + "task panicked during inline execution" + ); + + Err(DagError::TaskPanicked { + task_id: node_id.0, + panic_message, + }) + }), + ) + } else { + None + } + }) + .for_each(|fut| futures.push(fut)); + + while let Some(out) = futures.next().await { + if let Err(e) = out { + first_error.get_or_insert(e); } } - - // Drop the original sender so the channel closes when all tasks complete - drop(tx); - - // Wait for all tasks in this layer to complete - while rx.next().await.is_some() {} - } - } - - // Drop the output sender so the channel closes - drop(output_tx); - - // Collect all outputs and check for errors - let mut collected = Vec::new(); - let mut first_error = None; - while let Some((node_id, result)) = output_rx.next().await { - match result { - Ok(output) => collected.push((node_id, output)), - Err(e) if first_error.is_none() => first_error = Some(e), - Err(_) => {} // Ignore subsequent errors } } // Insert all successful outputs (avoiding holding lock across await) - let mut outputs = self.outputs.lock(); - for (node_id, output) in collected { - outputs.insert(node_id, output); + let mut outputs_lock = self.outputs.lock(); + for (node_id, output) in outputs { + outputs_lock.insert(node_id, output); } - drop(outputs); + drop(outputs_lock); // Return first error if any if let Some(err) = first_error { @@ -575,7 +547,7 @@ impl DagRunner { /// // Construct task with specific setting value /// let task = dag.add_task(Configuration::new(42)); /// - /// dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); + /// dag.run().await.unwrap(); /// /// assert_eq!(dag.get(task).unwrap(), 42); /// # }; diff --git a/src/task.rs b/src/task.rs index f362e87..54ee361 100644 --- a/src/task.rs +++ b/src/task.rs @@ -3,7 +3,7 @@ //! This module defines the core Task trait that all DAG nodes must implement, //! along with convenience functions for creating tasks from closures. -use std::future::Future; +use std::{future::Future, sync::Arc}; /// A unit of async work with typed inputs and outputs. /// @@ -106,7 +106,7 @@ pub trait Task: Send { #[doc(hidden)] fn extract_and_run( self, - receivers: Vec>, + dependencies: Vec>, ) -> impl Future> + Send; } @@ -149,10 +149,10 @@ where async fn extract_and_run( self, - receivers: Vec>, + dependencies: Vec>, ) -> Result { // For task_fn, we use ExtractInput since we don't have macro-generated extraction - let input = I::extract_from_channels(receivers) + let input = I::extract_from_deps(dependencies) .await .map_err(|e| format!("Failed to extract input: {}", e))?; diff --git a/src/types.rs b/src/types.rs index cca4e5d..7968d60 100644 --- a/src/types.rs +++ b/src/types.rs @@ -33,7 +33,7 @@ pub(crate) struct NodeId(pub usize); /// let dag = DagRunner::new(); /// let node = dag.add_task(LoadValue::new(42)); /// -/// dag.run(|fut| { tokio::spawn(fut); }).await.unwrap(); +/// dag.run().await.unwrap(); /// /// assert_eq!(dag.get(node).unwrap(), 42); /// # };