Wire up async runtimes

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

@@ -1,6 +1,7 @@
 use graph_craft::document::value::{RenderOutputType, TaggedValue, UVec2};
 use graph_craft::graphene_compiler::Executor;
 use graphene_std::application_io::{ExportFormat, RenderConfig, TimingInformation};
+use graphene_std::core_types::gpoll::GPoll;
 use graphene_std::core_types::ops::ConvertAsync;
 use graphene_std::core_types::transform::Footprint;
 use graphene_std::raster_types::{CPU, GPU, Raster};
@@ -8,8 +9,31 @@ use interpreted_executor::dynamic_executor::DynamicExecutor;
 use std::error::Error;
 use std::io::Cursor;
 use std::path::{Path, PathBuf};
+use std::sync::mpsc::Receiver;
 use std::time::Duration;
 
+const SOURCE_COMPLETION_TIMEOUT: Duration = Duration::from_secs(30);
+
+async fn execute_until_final(executor: &DynamicExecutor, render_config: RenderConfig, completion: &Receiver<()>) -> Result<TaggedValue, Box<dyn Error>> {
+	loop {
+		while completion.try_recv().is_ok() {}
+		match executor.execute(render_config.clone()).await? {
+			GPoll::Final(value) => return Ok(value),
+			GPoll::Fallback(boxed) => {
+				let (value, error) = *boxed;
+				log::error!("Node graph evaluation reported an error alongside its fallback output: {error:?}");
+				return Ok(value);
+			}
+			GPoll::Partial(_) | GPoll::Pending => {
+				completion
+					.recv_timeout(SOURCE_COMPLETION_TIMEOUT)
+					.map_err(|_| format!("Timed out after {}s waiting for async sources to complete", SOURCE_COMPLETION_TIMEOUT.as_secs()))?;
+			}
+			GPoll::Error(error) => return Err(format!("Node graph evaluation failed: {error:?}").into()),
+		}
+	}
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum FileType {
 	Svg,
@@ -36,6 +60,7 @@ pub async fn export_document(
 	scale: f64,
 	(width, height): (Option<u32>, Option<u32>),
 	transparent: bool,
+	completion: &Receiver<()>,
 ) -> Result<(), Box<dyn Error>> {
 	// Determine export format based on file type
 	let export_format = match file_type {
@@ -57,10 +82,7 @@ pub async fn export_document(
 	}
 
 	// Execute the graph
-	let (result, evaluation_error) = executor.execute(render_config).await?;
-	if let Some(error) = evaluation_error {
-		log::error!("Node graph evaluation reported an error alongside its fallback output: {error:?}");
-	}
+	let result = execute_until_final(executor, render_config, completion).await?;
 
 	// Handle the result based on output type
 	match result {
@@ -159,6 +181,7 @@ pub async fn export_gif(
 	scale: f64,
 	(width, height): (Option<u32>, Option<u32>),
 	animation: AnimationParams,
+	completion: &Receiver<()>,
 ) -> Result<(), Box<dyn Error>> {
 	use image::codecs::gif::{GifEncoder, Repeat};
 	use image::{Frame, RgbaImage};
@@ -198,10 +221,7 @@ pub async fn export_gif(
 		}
 
 		// Execute the graph for this frame
-		let (result, evaluation_error) = executor.execute(render_config).await?;
-		if let Some(error) = evaluation_error {
-			log::error!("Node graph evaluation reported an error alongside its fallback output: {error:?}");
-		}
+		let result = execute_until_final(executor, render_config, completion).await?;
 
 		// Extract RGBA data from result
 		let (data, img_width, img_height) = match result {

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

@@ -14,7 +14,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, NoopSpawner, RuntimeHandle};
+use graphene_std::runtime::{DynGraphRuntime, DynSpawner, GraphRuntime, RuntimeHandle, SourceFuture, Spawner};
 use interpreted_executor::dynamic_executor::DynamicExecutor;
 use interpreted_executor::util::wrap_network_in_scope;
 use std::error::Error;
@@ -29,6 +29,29 @@ impl NodeGraphUpdateSender for UpdateLogger {
 	}
 }
 
+struct TokioSpawner(Option<tokio::runtime::Runtime>);
+
+impl TokioSpawner {
+	fn new() -> Self {
+		Self(Some(tokio::runtime::Runtime::new().expect("Failed to start the async source runtime")))
+	}
+}
+
+impl Spawner for TokioSpawner {
+	fn spawn(&self, task: SourceFuture) {
+		self.0.as_ref().expect("runtime lives until drop").spawn(task);
+	}
+}
+
+/// Dropping a tokio runtime blocks on its tasks, which panics inside the async main; shut down in the background instead.
+impl Drop for TokioSpawner {
+	fn drop(&mut self) {
+		if let Some(runtime) = self.0.take() {
+			runtime.shutdown_background();
+		}
+	}
+}
+
 #[derive(Debug, Parser)]
 #[clap(name = "graphene-cli", version)]
 pub struct App {
@@ -183,7 +206,11 @@ async fn main() -> Result<(), Box<dyn Error>> {
 	let preferences = EditorPreferences {
 		max_render_region_size: EditorPreferences::default().max_render_region_size,
 	};
-	let graph_runtime: Arc<DynGraphRuntime> = Arc::new(GraphRuntime::new(Box::new(NoopSpawner) as Box<DynSpawner>));
+	let graph_runtime: Arc<DynGraphRuntime> = Arc::new(GraphRuntime::new(Box::new(TokioSpawner::new()) as Box<DynSpawner>));
+	let (completion_sender, completion_receiver) = std::sync::mpsc::channel();
+	graph_runtime.set_notifier(Arc::new(move || {
+		let _ = completion_sender.send(());
+	}));
 	let editor_api = Arc::new(PlatformEditorApi {
 		application_io: Some(application_io_for_api),
 		node_graph_message_sender: Box::new(UpdateLogger {}),
@@ -230,9 +257,9 @@ async fn main() -> Result<(), Box<dyn Error>> {
 			// Perform export based on file type
 			if file_type == export::FileType::Gif {
 				let animation = export::AnimationParams::new(fps, frames, duration);
-				export::export_gif(&executor, wgpu_executor_ref.clone(), output, scale, (width, height), animation).await?;
+				export::export_gif(&executor, wgpu_executor_ref.clone(), output, scale, (width, height), animation, &completion_receiver).await?;
 			} else {
-				export::export_document(&executor, wgpu_executor_ref.clone(), output, file_type, scale, (width, height), transparent).await?;
+				export::export_document(&executor, wgpu_executor_ref.clone(), output, file_type, scale, (width, height), transparent, &completion_receiver).await?;
 			}
 		}
 		_ => unreachable!("All other commands should be handled before this match statement is run"),