Poll source tasks once at spawn and land inline completions in the first eval

node-graph/graphene-cli/src/main.rs

@@ -13,7 +13,7 @@ use graph_craft::graphene_compiler::Compiler;
 use graph_craft::proto::ProtoNetwork;
 use graph_craft::util::load_network;
 use graphene_std::application_io::{ApplicationIo, NodeGraphUpdateMessage, NodeGraphUpdateSender};
-use graphene_std::runtime::{DynGraphRuntime, DynSpawner, GraphRuntime, RuntimeHandle, SourceFuture, Spawner};
+use graphene_std::runtime::{DynGraphRuntime, DynSpawner, GraphRuntime, RuntimeHandle, SourceFuture, Spawner, poll_once};
 use interpreted_executor::dynamic_executor::DynamicExecutor;
 use interpreted_executor::util::wrap_network_in_scope;
 use std::error::Error;
@@ -37,8 +37,14 @@ impl TokioSpawner {
 }
 
 impl Spawner for TokioSpawner {
-	fn spawn(&self, task: SourceFuture) {
-		self.0.as_ref().expect("runtime lives until drop").spawn(task);
+	fn spawn(&self, mut task: SourceFuture) -> bool {
+		let runtime = self.0.as_ref().expect("runtime lives until drop");
+		let _guard = runtime.enter();
+		if poll_once(&mut task) {
+			return true;
+		}
+		runtime.spawn(task);
+		false
 	}
 }
 

node-graph/interpreted-executor/src/dynamic_executor.rs

@@ -482,7 +482,9 @@ mod test {
 	struct InertSpawner;
 
 	impl Spawner for InertSpawner {
-		fn spawn(&self, _task: SourceFuture) {}
+		fn spawn(&self, _task: SourceFuture) -> bool {
+			false
+		}
 	}
 
 	#[test]

node-graph/libraries/core-types/src/runtime.rs

@@ -16,7 +16,8 @@ pub type DynRuntime = dyn Runtime + Send + Sync;
 pub type DynRuntime = dyn Runtime;
 
 pub trait Runtime {
-	fn spawn(&self, source: SourceId, future: SourceFuture);
+	/// Returns true when the future completed during the call, so its result is already observable.
+	fn spawn(&self, source: SourceId, future: SourceFuture) -> bool;
 }
 
 #[derive(Clone)]
@@ -40,7 +41,14 @@ impl graphene_hash::CacheHash for RuntimeHandle {
 }
 
 pub trait Spawner {
-	fn spawn(&self, task: SourceFuture);
+	/// Returns true when the task completed during the call, so its result is already observable.
+	fn spawn(&self, task: SourceFuture) -> bool;
+}
+
+/// Polls `task` once with a no-op waker, returning true if it completed.
+pub fn poll_once(task: &mut SourceFuture) -> bool {
+	let mut context = std::task::Context::from_waker(std::task::Waker::noop());
+	task.as_mut().poll(&mut context).is_ready()
 }
 
 #[cfg(not(target_family = "wasm"))]
@@ -54,7 +62,7 @@ pub type DynNotifier = dyn Fn() + Send + Sync;
 pub type DynNotifier = dyn Fn();
 
 impl<S: Spawner + ?Sized> Spawner for Box<S> {
-	fn spawn(&self, task: SourceFuture) {
+	fn spawn(&self, task: SourceFuture) -> bool {
 		(**self).spawn(task)
 	}
 }
@@ -63,11 +71,12 @@ impl<S: Spawner + ?Sized> Spawner for Box<S> {
 pub struct NoopSpawner;
 
 impl Spawner for NoopSpawner {
-	fn spawn(&self, mut task: SourceFuture) {
-		let mut context = std::task::Context::from_waker(std::task::Waker::noop());
-		if task.as_mut().poll(&mut context).is_pending() {
-			log::warn!("async source is not immediately ready and no host spawner is wired; the task is dropped");
+	fn spawn(&self, mut task: SourceFuture) -> bool {
+		if poll_once(&mut task) {
+			return true;
 		}
+		log::warn!("async source is not immediately ready and no host spawner is wired; the task is dropped");
+		false
 	}
 }
 
@@ -132,21 +141,26 @@ impl<S> GraphRuntime<S> {
 }
 
 impl<S: Spawner> Runtime for GraphRuntime<S> {
-	fn spawn(&self, source: SourceId, future: SourceFuture) {
+	fn spawn(&self, source: SourceId, mut future: SourceFuture) -> bool {
 		let generations = Arc::clone(&self.generations);
 		let dirty = Arc::clone(&self.dirty);
 		let notifier = Arc::clone(&self.notifier);
-		self.spawner.spawn(Box::pin(async move {
-			future.await;
-			let mut generations = generations.lock().unwrap_or_else(PoisonError::into_inner);
-			if let Some(generation) = generations.get_mut(&source) {
-				*generation += 1;
-				dirty.store(true, Ordering::Release);
-				drop(generations);
-				let notifier = Arc::clone(&notifier.lock().unwrap_or_else(PoisonError::into_inner));
-				notifier();
+		let mut first = true;
+		self.spawner.spawn(Box::pin(std::future::poll_fn(move |task_context| {
+			let poll = future.as_mut().poll(task_context);
+			if poll.is_ready() && !first {
+				let mut generations = generations.lock().unwrap_or_else(PoisonError::into_inner);
+				if let Some(generation) = generations.get_mut(&source) {
+					*generation += 1;
+					dirty.store(true, Ordering::Release);
+					drop(generations);
+					let notifier = Arc::clone(&notifier.lock().unwrap_or_else(PoisonError::into_inner));
+					notifier();
+				}
 			}
-		}));
+			first = false;
+			poll
+		})))
 	}
 }
 
@@ -167,8 +181,9 @@ mod tests {
 	}
 
 	impl Runtime for MockRuntime {
-		fn spawn(&self, source: SourceId, future: SourceFuture) {
+		fn spawn(&self, source: SourceId, future: SourceFuture) -> bool {
 			self.futures.lock().unwrap().push((source, future));
+			false
 		}
 	}
 
@@ -192,11 +207,34 @@ mod tests {
 	}
 
 	impl Spawner for CollectSpawner {
-		fn spawn(&self, task: SourceFuture) {
+		fn spawn(&self, mut task: SourceFuture) -> bool {
+			if poll_once(&mut task) {
+				return true;
+			}
 			self.tasks.lock().unwrap().push(task);
+			false
 		}
 	}
 
+	struct YieldOnce(bool);
+
+	impl Future for YieldOnce {
+		type Output = ();
+
+		fn poll(mut self: Pin<&mut Self>, task_context: &mut std::task::Context<'_>) -> std::task::Poll<()> {
+			if self.0 {
+				return std::task::Poll::Ready(());
+			}
+			self.0 = true;
+			task_context.waker().wake_by_ref();
+			std::task::Poll::Pending
+		}
+	}
+
+	fn yield_once() -> YieldOnce {
+		YieldOnce(false)
+	}
+
 	impl CollectSpawner {
 		fn drain(&self) -> usize {
 			let tasks = std::mem::take(&mut *self.tasks.lock().unwrap());
@@ -413,7 +451,7 @@ mod tests {
 		let runtime = GraphRuntime::new(CollectSpawner::default());
 		runtime.retain_sources(&[7]);
 
-		Runtime::spawn(&runtime, 7, Box::pin(async {}));
+		Runtime::spawn(&runtime, 7, Box::pin(yield_once()));
 		assert_eq!(runtime.snapshot(), vec![(7, 0)], "no bump before the future completes");
 		assert!(!runtime.take_dirty());
 
@@ -434,7 +472,7 @@ mod tests {
 			observed.store(dirty_at_notify.load(Ordering::Acquire), Ordering::Relaxed);
 		}));
 
-		Runtime::spawn(&runtime, 7, Box::pin(async {}));
+		Runtime::spawn(&runtime, 7, Box::pin(yield_once()));
 		assert_eq!(runtime.spawner().drain(), 1);
 		assert!(observed_dirty.load(Ordering::Relaxed), "the notifier must observe the dirty flag already set");
 	}
@@ -447,7 +485,7 @@ mod tests {
 		let flag = Arc::clone(&notified);
 		runtime.set_notifier(Arc::new(move || flag.store(true, Ordering::Relaxed)));
 
-		Runtime::spawn(&runtime, 7, Box::pin(async {}));
+		Runtime::spawn(&runtime, 7, Box::pin(yield_once()));
 		runtime.retain_sources(&[]);
 		assert_eq!(runtime.spawner().drain(), 1);
 		assert!(!notified.load(Ordering::Relaxed));
@@ -458,7 +496,7 @@ mod tests {
 		let runtime = GraphRuntime::new(CollectSpawner::default());
 		runtime.retain_sources(&[7]);
 
-		Runtime::spawn(&runtime, 7, Box::pin(async {}));
+		Runtime::spawn(&runtime, 7, Box::pin(yield_once()));
 		runtime.retain_sources(&[]);
 		assert_eq!(runtime.spawner().drain(), 1);
 
@@ -470,7 +508,7 @@ mod tests {
 	fn retain_sources_preserves_live_generations() {
 		let runtime = GraphRuntime::new(CollectSpawner::default());
 		runtime.retain_sources(&[7]);
-		Runtime::spawn(&runtime, 7, Box::pin(async {}));
+		Runtime::spawn(&runtime, 7, Box::pin(yield_once()));
 		runtime.spawner().drain();
 
 		runtime.retain_sources(&[7, 9]);
@@ -482,9 +520,42 @@ mod tests {
 
 	#[node_macro::node(category(""))]
 	async fn epilogue_double(_: impl Ctx, value: f64) -> f64 {
+		yield_once().await;
 		value * 2.
 	}
 
+	#[node_macro::node(category(""))]
+	async fn inline_double(_: impl Ctx, value: f64) -> f64 {
+		value * 2.
+	}
+
+	#[test]
+	fn an_immediately_ready_task_completes_inline_without_invalidating() {
+		let runtime = GraphRuntime::new(CollectSpawner::default());
+		runtime.retain_sources(&[7]);
+
+		assert!(Runtime::spawn(&runtime, 7, Box::pin(async {})));
+		assert_eq!(runtime.spawner().drain(), 0);
+		assert_eq!(runtime.snapshot(), vec![(7, 0)], "inline completion must not bump the generation");
+		assert!(!runtime.take_dirty());
+	}
+
+	#[test]
+	fn an_immediately_ready_kernel_returns_final_on_the_first_eval() {
+		let arena = Arena::new(64).unwrap();
+		let runtime = Arc::new(GraphRuntime::new(CollectSpawner::default()));
+		runtime.retain_sources(&[13]);
+		let graph = InlineDoubleNode::new(SourceNode(21.0f64), SourceNode(RuntimeHandle(runtime.clone())), SourceNode(13u64));
+
+		let snapshot = runtime.snapshot();
+		let scope = EvalScope::new(None, None, None, &snapshot, &arena);
+		let ctx = ContextImpl::root(&scope);
+		assert_eq!(Node::eval(&graph, &ctx), GPoll::Final(42.0));
+		assert!(!runtime.take_dirty());
+		assert_eq!(runtime.snapshot(), vec![(13, 0)]);
+		assert_eq!(runtime.spawner().drain(), 0);
+	}
+
 	#[test]
 	fn a_source_slot_lands_through_the_runtime_while_downstream_keys_invalidate() {
 		let arena = Arena::new(64).unwrap();

node-graph/node-macro/src/codegen.rs

@@ -932,6 +932,22 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
 		},
 		_ => quote!(#core_types::gpoll::GPoll::Final(__future.await)),
 	};
+	let spawn_tail = |completion: TokenStream2, fallback: TokenStream2| {
+		quote! {
+			self.slot.lock().unwrap_or_else(std::sync::PoisonError::into_inner).insert(__key, None);
+			let __slot = std::sync::Arc::clone(&self.slot);
+			if _runtime.0.spawn(_source, Box::pin(async move {
+				let __value = #completion;
+				__slot.lock().unwrap_or_else(std::sync::PoisonError::into_inner).insert(__key, Some(__value));
+			})) {
+				let __entries = self.slot.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
+				if let Some(Some(__value)) = __entries.get(&__key) {
+					return __cell.merge(__value.clone());
+				}
+			}
+			#fallback
+		}
+	};
 	let eval_tail = match (async_fn, future_kernel) {
 		(false, false) => lift,
 		(true, _) => {
@@ -942,17 +958,12 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
 				.chain(data_names.iter().map(|name| quote!(self.#name.clone())))
 				.chain(kernel_value_names.iter().map(|name| quote!(#name.clone())));
 			let completion = future_completion(&parsed.output_type);
+			let tail = spawn_tail(completion, inflight.clone());
 			quote! {
 				#slot_check
-				self.slot.lock().unwrap_or_else(std::sync::PoisonError::into_inner).insert(__key, None);
-				let __slot = std::sync::Arc::clone(&self.slot);
 				#(#snapshot_binding)*
 				let __future = self::#fn_name(#(#future_args),*);
-				_runtime.0.spawn(_source, Box::pin(async move {
-					let __value = #completion;
-					__slot.lock().unwrap_or_else(std::sync::PoisonError::into_inner).insert(__key, Some(__value));
-				}));
-				#inflight
+				#tail
 			}
 		}
 		(false, true) => {
@@ -977,17 +988,12 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
 				_ => unreachable!("guarded by future_kernel"),
 			};
 			let completion = future_completion(&payload);
+			let tail = spawn_tail(completion, spawn_return);
 			quote! {
 				#slot_check
 				#placeholder_binding
 				#acquire
-				self.slot.lock().unwrap_or_else(std::sync::PoisonError::into_inner).insert(__key, None);
-				let __slot = std::sync::Arc::clone(&self.slot);
-				_runtime.0.spawn(_source, Box::pin(async move {
-					let __value = #completion;
-					__slot.lock().unwrap_or_else(std::sync::PoisonError::into_inner).insert(__key, Some(__value));
-				}));
-				#spawn_return
+				#tail
 			}
 		}
 	};