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

Don't use channels to send dependencies

Browse Source
This commit is contained in:
TechnoPorg committed 2026-01-19 00:24:32 +01:00
1 parent 8e507d76f3
commit df39dd80b9
14 files changed
+212 -300
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Cargo.lock
Generated
+21
View File
@@ -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"
Cargo.toml
+11 -2
View File
@@ -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"
README.md
+5 -5
View File
@@ -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
dagx-macros/src/lib.rs
+16 -35
View File
@@ -335,7 +335,7 @@ pub fn task(_attr: TokenStream, item: TokenStream) -> TokenStream {
quote! {
fn extract_and_run(
self,
_receivers: Vec<Box<dyn std::any::Any + Send>>,
_dependencies: Vec<std::sync::Arc<dyn std::any::Any + Send + Sync>>,
) -> impl std::future::Future<Output = Result<Self::Output, String>> + 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<Box<dyn std::any::Any + Send>>,
mut dependencies: Vec<std::sync::Arc<dyn std::any::Any + Send + Sync>>,
) -> impl std::future::Future<Output = Result<Self::Output, String>> + 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::<oneshot::Receiver<Arc<#param_type>>>()
.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<Box<dyn std::any::Any + Send>>,
dependencies: Vec<std::sync::Arc<dyn std::any::Any + Send + Sync>>,
) -> impl std::future::Future<Output = Result<Self::Output, String>> + 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::<oneshot::Receiver<Arc<#param_types>>>()
.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)
docs/COMPARISONS.md
+1 -1
View File
@@ -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):
docs/CYCLE_PREVENTION.md
+1 -1
View File
@@ -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:**
src/builder.rs
+2 -2
View File
@@ -65,7 +65,7 @@ impl IsUnitType for () {}
/// let b = b.depends_on(&a);
/// // Now b is a TaskHandle<i32>
///
/// 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<D>(self, deps: D) -> TaskHandle<Tk::Output>
src/error.rs
+1 -1
View File
@@ -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
)
}
src/extract.rs
+50 -79
View File
@@ -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<Box<dyn Any + Send>>,
fn extract_from_deps(
deps: Vec<Arc<dyn Any + Send + Sync>>,
) -> impl std::future::Future<Output = Result<Self, String>> + Send;
}
// Unit type - no dependencies
impl ExtractInput for () {
async fn extract_from_channels(_receivers: Vec<Box<dyn Any + Send>>) -> Result<Self, String> {
async fn extract_from_deps(_deps: Vec<Arc<dyn Any + Send + Sync>>) -> Result<Self, String> {
Ok(())
}
}
@@ -67,15 +65,14 @@ macro_rules! impl_extract_single {
($($t:ty),+) => {
$(
impl ExtractInput for $t {
async fn extract_from_channels(mut receivers: Vec<Box<dyn Any + Send>>) -> Result<Self, String> {
if receivers.len() != 1 {
return Err(format!("Expected 1 dependency, got {}", receivers.len()));
async fn extract_from_deps(mut dependencies: Vec<Arc<dyn Any + Send + Sync>>) -> Result<Self, String> {
if dependencies.len() != 1 {
return Err(format!("Expected 1 dependency, got {}", dependencies.len()));
}
let rx = *receivers.pop()
let arc_value = dependencies.pop()
.unwrap()
.downcast::<oneshot::Receiver<Arc<$t>>>()
.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<Box<dyn Any + Send>>,
async fn extract_from_deps(
mut dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> Result<Self, String> {
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::<oneshot::Receiver<Arc<Result<T, E>>>>()
.downcast::<Result<T, E>>()
.map_err(|_| {
format!(
"Type mismatch: expected Arc<Result<{}, {}>>",
@@ -110,9 +107,6 @@ where
std::any::type_name::<E>()
)
})?;
let arc_value = rx
.await
.map_err(|_| "Channel closed before receiving value".to_string())?;
Ok((*arc_value).clone())
}
}
@@ -122,25 +116,22 @@ impl<T> ExtractInput for Option<T>
where
T: 'static + Clone + Send + Sync,
{
async fn extract_from_channels(
mut receivers: Vec<Box<dyn Any + Send>>,
async fn extract_from_deps(
mut dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> Result<Self, String> {
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::<oneshot::Receiver<Arc<Option<T>>>>()
.downcast::<Option<T>>()
.map_err(|_| {
format!(
"Type mismatch: expected Arc<Option<{}>>",
std::any::type_name::<T>()
)
})?;
let arc_value = rx
.await
.map_err(|_| "Channel closed before receiving value".to_string())?;
Ok((*arc_value).clone())
}
}
@@ -150,25 +141,22 @@ impl<T> ExtractInput for Vec<T>
where
T: 'static + Clone + Send + Sync,
{
async fn extract_from_channels(
mut receivers: Vec<Box<dyn Any + Send>>,
async fn extract_from_deps(
mut dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> Result<Self, String> {
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::<oneshot::Receiver<Arc<Vec<T>>>>()
.downcast::<Vec<T>>()
.map_err(|_| {
format!(
"Type mismatch: expected Arc<Vec<{}>>",
std::any::type_name::<T>()
)
})?;
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<Box<dyn Any + Send>>,
async fn extract_from_deps(
mut dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> Result<Self, String> {
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::<oneshot::Receiver<Arc<std::collections::HashMap<K, V>>>>()
.downcast::<std::collections::HashMap<K, V>>()
.map_err(|_| {
format!(
"Type mismatch: expected Arc<HashMap<{}, {}>>",
@@ -196,9 +184,6 @@ where
std::any::type_name::<V>()
)
})?;
let arc_value = rx
.await
.map_err(|_| "Channel closed before receiving value".to_string())?;
Ok((*arc_value).clone())
}
}
@@ -209,25 +194,22 @@ impl<T> ExtractInput for Arc<T>
where
T: 'static + Clone + Send + Sync,
{
async fn extract_from_channels(
mut receivers: Vec<Box<dyn Any + Send>>,
async fn extract_from_deps(
mut dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> Result<Self, String> {
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::<oneshot::Receiver<Arc<Arc<T>>>>()
.downcast::<Arc<T>>()
.map_err(|_| {
format!(
"Type mismatch: expected Arc<Arc<{}>>",
std::any::type_name::<T>()
)
})?;
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<Box<dyn Any + Send>>) -> Result<Self, String> {
// Verify we have the correct number of receivers
async fn extract_from_deps(dependencies: Vec<Arc<dyn Any + Send + Sync>>) -> Result<Self, String> {
// 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
let arc_results = (
$(
#[allow(non_snake_case)]
let $T = *iter.next()
iter.next()
.ok_or_else(|| format!("Missing receiver at index {}", $idx))?
.downcast::<oneshot::Receiver<Arc<$T>>>()
.downcast::<$T>()
.map_err(|_| format!("Type mismatch at index {}: expected Arc<{}>",
$idx, std::any::type_name::<$T>()))?;
$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!(
$(
async move {
$T.await.map_err(|_| format!("Channel {} closed", $idx))
}
),+
);
// Clone inner values from Arc (cheap for small types, necessary for ownership)
Ok(($(
(*arc_results.$idx?).clone()
(*arc_results.$idx).clone()
,)+))
}
}
src/lib.rs
+8 -8
View File
@@ -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();
//! # };
//! ```
//!
src/node.rs
+5 -47
View File
@@ -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<Box<dyn Any + Send>>;
/// Type-erased channel pair (senders, receivers)
type ChannelPair = (ErasedChannels, ErasedChannels);
/// Type-erased async execution result (Arc-wrapped for efficient fanout)
type ExecuteFuture =
Pin<Box<dyn Future<Output = Result<Arc<dyn Any + Send + Sync>, 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<Self>,
receivers: ErasedChannels,
senders: ErasedChannels,
dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> 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::<Arc<T::Output>>();
senders.push(Box::new(tx) as Box<dyn Any + Send>);
receivers.push(Box::new(rx) as Box<dyn Any + Send>);
}
(senders, receivers)
}
fn execute_with_channels(
fn execute_with_deps(
self: Box<Self>,
receivers: ErasedChannels,
senders: ErasedChannels,
dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> 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<T> knows T::Output, so downcast is safe
for sender_any in senders {
if let Ok(sender) = sender_any.downcast::<oneshot::Sender<Arc<T::Output>>>() {
// 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<dyn Any + Send + Sync>)
})
src/runner.rs
+86 -114
View File
@@ -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<Vec<Option<Box<dyn ExecutableNode>>>>,
pub(crate) nodes: Mutex<Vec<Option<Box<dyn ExecutableNode + Sync>>>>,
pub(crate) outputs: Mutex<HashMap<NodeId, std::sync::Arc<dyn std::any::Any + Send + Sync>>>,
pub(crate) edges: Mutex<HashMap<NodeId, Vec<NodeId>>>, // node -> dependencies
pub(crate) dependents: Mutex<HashMap<NodeId, Vec<NodeId>>>, // 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<Tk>(&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<S>(&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<NodeId, Vec<Box<dyn std::any::Any + Send>>> =
HashMap::new();
let mut producer_senders: HashMap<NodeId, Vec<Box<dyn std::any::Any + Send>>> =
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<std::sync::Arc<dyn std::any::Any + Send + Sync>, DagError>,
)>();
let outputs: DashMap<NodeId, Arc<dyn Any + Send + Sync>> = 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 {
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");
let mut task_tx = tx.clone();
let out_tx = output_tx.clone();
// 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(());
};
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));
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::<String>()
{
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);
/// # };
src/task.rs
+4 -4
View File
@@ -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<Box<dyn std::any::Any + Send>>,
dependencies: Vec<Arc<dyn std::any::Any + Send + Sync>>,
) -> impl Future<Output = Result<Self::Output, String>> + Send;
}
@@ -149,10 +149,10 @@ where
async fn extract_and_run(
self,
receivers: Vec<Box<dyn std::any::Any + Send>>,
dependencies: Vec<Arc<dyn std::any::Any + Send + Sync>>,
) -> Result<Self::Output, String> {
// 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))?;
src/types.rs
+1 -1
View File
@@ -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);
/// # };