Implement editor level future message handler (#4182)

* Implement proper editor level async message handler

* Take application_io and wake callback as Editor::new arguments

* Rename AsyncMessage -> FutureMessage

* add DesktopWrapperMessage::Wake

---------

Co-authored-by: Timon <me@timon.zip>

editor/Cargo.toml

@@ -55,8 +55,13 @@ wgpu-executor = { workspace = true, optional = true }
 # Optional workspace dependencies
 wasm-bindgen = { workspace = true, optional = true }
 
+[target.'cfg(not(target_family = "wasm"))'.dependencies]
+tokio = { workspace = true }
+
+[target.'cfg(target_family = "wasm")'.dependencies]
+wasm-bindgen-futures = { workspace = true }
+
 [dev-dependencies]
 # Workspace dependencies
 env_logger = { workspace = true }
 futures = { workspace = true }
-tokio = { workspace = true }

editor/src/application.rs

@@ -10,13 +10,16 @@ pub struct Editor {
 }
 
 impl Editor {
-	pub fn new(environment: Environment, uuid_random_seed: u64, resource_storage: Box<dyn ResourceStorage>) -> Self {
+	pub fn new(environment: Environment, uuid_random_seed: u64, resource_storage: Box<dyn ResourceStorage>, mut application_io: PlatformApplicationIo, wake: Wake) -> Self {
 		ENVIRONMENT.set(environment).expect("Editor shoud only be initialized once");
 		graphene_std::uuid::set_uuid_seed(uuid_random_seed);
 
-		Self {
-			dispatcher: Dispatcher::new(resource_storage),
-		}
+		let mut dispatcher = Dispatcher::new(resource_storage);
+		dispatcher.message_handlers.async_message_handler.set_wake(wake);
+		application_io.inject_resource_proxy(dispatcher.message_handlers.resource_storage_message_handler.resources());
+		crate::node_graph_executor::replace_application_io(application_io);
+
+		Self { dispatcher }
 	}
 
 	#[cfg(test)]
@@ -45,11 +48,6 @@ impl Editor {
 	pub fn poll_node_graph_evaluation(&mut self, responses: &mut VecDeque<Message>) -> Result<(), String> {
 		self.dispatcher.poll_node_graph_evaluation(responses)
 	}
-
-	pub fn replace_application_io(&mut self, mut application_io: PlatformApplicationIo) {
-		application_io.inject_resource_proxy(self.dispatcher.message_handlers.resource_storage_message_handler.resources());
-		crate::node_graph_executor::replace_application_io(application_io)
-	}
 }
 
 static ENVIRONMENT: OnceLock<Environment> = OnceLock::new();

editor/src/dispatcher.rs

@@ -20,6 +20,7 @@ pub struct Dispatcher {
 pub struct DispatcherMessageHandlers {
 	animation_message_handler: AnimationMessageHandler,
 	app_window_message_handler: AppWindowMessageHandler,
+	pub(crate) async_message_handler: FutureMessageHandler,
 	broadcast_message_handler: BroadcastMessageHandler,
 	clipboard_message_handler: ClipboardMessageHandler,
 	color_picker_message_handler: ColorPickerMessageHandler,
@@ -124,6 +125,13 @@ impl Dispatcher {
 	pub fn handle_message<T: Into<Message>>(&mut self, message: T, process_after_all_current: bool) {
 		let message = message.into();
 
+		// Drain async results into the queue before processing the new message.
+		let mut async_results = VecDeque::new();
+		self.message_handlers.async_message_handler.drain_results(&mut async_results);
+		if !async_results.is_empty() {
+			Self::schedule_execution(&mut self.message_queues, true, async_results);
+		}
+
 		// If we are not maintaining the buffer, simply add to the current queue
 		Self::schedule_execution(&mut self.message_queues, process_after_all_current, [message]);
 
@@ -173,6 +181,9 @@ impl Dispatcher {
 				Message::AppWindow(message) => {
 					self.message_handlers.app_window_message_handler.process_message(message, &mut queue, ());
 				}
+				Message::Future(message) => {
+					self.message_handlers.async_message_handler.process_message(message, &mut queue, FutureMessageContext {});
+				}
 				Message::Broadcast(message) => self.message_handlers.broadcast_message_handler.process_message(message, &mut queue, ()),
 				Message::Clipboard(message) => self.message_handlers.clipboard_message_handler.process_message(message, &mut queue, ()),
 				Message::ColorPicker(message) => self.message_handlers.color_picker_message_handler.process_message(message, &mut queue, ()),

editor/src/messages/frontend/frontend_message.rs

@@ -1,7 +1,7 @@
 use super::IconName;
 use super::utility_types::{MouseCursorIcon, PersistedState};
 use crate::messages::app_window::app_window_message_handler::AppWindowPlatform;
-use crate::messages::frontend::utility_types::{DocumentInfo, EyedropperPreviewImage, FrontendMessageFuture};
+use crate::messages::frontend::utility_types::{DocumentInfo, EyedropperPreviewImage};
 use crate::messages::input_mapper::utility_types::misc::ActionShortcut;
 use crate::messages::layout::utility_types::widget_prelude::*;
 use crate::messages::portfolio::document::node_graph::utility_types::{
@@ -27,13 +27,6 @@ use crate::messages::portfolio::document::overlays::utility_types::OverlayContex
 #[derive(derivative::Derivative, Clone, serde::Serialize, serde::Deserialize)]
 #[derivative(Debug, PartialEq)]
 pub enum FrontendMessage {
-	Await {
-		#[serde(skip, default)]
-		#[derivative(Debug = "ignore", PartialEq = "ignore")]
-		#[cfg_attr(feature = "wasm", tsify(type = "unknown"))]
-		future: FrontendMessageFuture,
-	},
-
 	// Display prefix: make the frontend show something, like a dialog
 	DisplayDialog {
 		title: String,

editor/src/messages/frontend/mod.rs

@@ -4,7 +4,6 @@ pub mod utility_types;
 
 #[doc(inline)]
 pub use frontend_message::{FrontendMessage, FrontendMessageDiscriminant};
-pub use utility_types::FrontendMessageFuture;
 
 // TODO: Make this an enum with the actual icon names, somehow derived from or tied to the frontend icon set.
 // TODO: Then remove `#[widget_builder(string)]` from all icon fields.

editor/src/messages/frontend/utility_types.rs

@@ -1,7 +1,4 @@
-use std::future::{Future, IntoFuture};
 use std::path::PathBuf;
-use std::pin::Pin;
-use std::sync::{Arc, Mutex};
 
 use graph_craft::application_io::resource::ResourceHash;
 
@@ -85,31 +82,3 @@ pub struct EyedropperPreviewImage {
 	pub width: u32,
 	pub height: u32,
 }
-
-#[derive(Clone, Default)]
-pub struct FrontendMessageFuture {
-	inner: Arc<Mutex<Option<InnerFrontendMessageFuture>>>,
-}
-
-impl FrontendMessageFuture {
-	pub fn new(future: impl Future<Output = FrontendMessage> + Send + 'static) -> Self {
-		Self {
-			inner: Arc::new(Mutex::new(Some(Box::pin(future)))),
-		}
-	}
-}
-
-type InnerFrontendMessageFuture = Pin<Box<dyn Future<Output = FrontendMessage> + Send + 'static>>;
-
-impl IntoFuture for FrontendMessageFuture {
-	type Output = FrontendMessage;
-	type IntoFuture = InnerFrontendMessageFuture;
-
-	fn into_future(self) -> Self::IntoFuture {
-		self.inner
-			.lock()
-			.unwrap_or_else(|poisoned| poisoned.into_inner())
-			.take()
-			.expect("FrontendMessageFuture can only be awaited once")
-	}
-}

editor/src/messages/future/future_message.rs

@@ -0,0 +1,16 @@
+use crate::messages::future::MessageFuture;
+use crate::messages::prelude::*;
+
+#[impl_message(Message, Future)]
+#[derive(derivative::Derivative, Clone, serde::Serialize, serde::Deserialize)]
+#[derivative(Debug, PartialEq)]
+pub enum FutureMessage {
+	/// Spawn `future`; its resolved [`Message`] re-enters the dispatcher on the next tick.
+	Await {
+		#[serde(skip, default)]
+		#[derivative(Debug = "ignore", PartialEq = "ignore")]
+		future: MessageFuture,
+	},
+	/// Sent by a wake callback to nudge an idle event loop into a dispatch tick.
+	Wake,
+}

editor/src/messages/future/future_message_handler.rs

@@ -0,0 +1,174 @@
+use std::future::{Future, IntoFuture};
+use std::pin::Pin;
+use std::sync::{Arc, Mutex};
+
+use futures::channel::mpsc::{UnboundedReceiver, UnboundedSender, unbounded};
+
+use crate::messages::prelude::*;
+
+type InnerMessageFuture = Pin<Box<dyn Future<Output = Message> + Send + 'static>>;
+
+/// Invoked by the spawner after a result is sent, to wake the platform event loop.
+pub type Wake = Arc<dyn Fn() + Send + Sync>;
+
+fn noop_wake() -> Wake {
+	Arc::new(|| {})
+}
+
+/// One-shot async work whose result re-enters the dispatcher as a [`Message`].
+/// Resolves to [`Message::NoOp`] if polled after the inner future has already been taken.
+#[derive(Clone, Default)]
+pub struct MessageFuture {
+	inner: Arc<Mutex<Option<InnerMessageFuture>>>,
+}
+
+impl MessageFuture {
+	pub fn new(future: impl Future<Output = Message> + Send + 'static) -> Self {
+		Self {
+			inner: Arc::new(Mutex::new(Some(Box::pin(future)))),
+		}
+	}
+}
+
+impl IntoFuture for MessageFuture {
+	type Output = Message;
+	type IntoFuture = InnerMessageFuture;
+
+	fn into_future(self) -> Self::IntoFuture {
+		let taken = self.inner.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).take();
+		match taken {
+			Some(future) => future,
+			None => Box::pin(async { Message::NoOp }),
+		}
+	}
+}
+
+/// Platform-specific async-task executor.
+/// Runs `future`, sends the resolved message on `results`, then calls `wake`.
+pub trait MessageSpawner: Send + Sync {
+	fn spawn(&self, future: InnerMessageFuture, results: UnboundedSender<Message>, wake: Wake);
+}
+
+#[derive(ExtractField)]
+pub struct FutureMessageContext {}
+
+#[derive(ExtractField)]
+pub struct FutureMessageHandler {
+	spawner: Arc<dyn MessageSpawner>,
+	wake: Wake,
+	results_sender: UnboundedSender<Message>,
+	results_receiver: UnboundedReceiver<Message>,
+}
+
+impl FutureMessageHandler {
+	pub fn with_wake(wake: Wake) -> Self {
+		let (results_sender, results_receiver) = unbounded();
+		Self {
+			spawner: default_spawner(),
+			wake,
+			results_sender,
+			results_receiver,
+		}
+	}
+
+	pub fn set_wake(&mut self, wake: Wake) {
+		self.wake = wake;
+	}
+
+	/// Pull every resolved async result into `out`.
+	pub fn drain_results(&mut self, out: &mut VecDeque<Message>) {
+		while let Ok(Some(message)) = self.results_receiver.try_next() {
+			out.push_back(message);
+		}
+	}
+}
+
+impl Default for FutureMessageHandler {
+	fn default() -> Self {
+		Self::with_wake(noop_wake())
+	}
+}
+
+impl std::fmt::Debug for FutureMessageHandler {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		f.debug_struct("FutureMessageHandler").finish_non_exhaustive()
+	}
+}
+
+#[message_handler_data]
+impl MessageHandler<FutureMessage, FutureMessageContext> for FutureMessageHandler {
+	fn process_message(&mut self, message: FutureMessage, _responses: &mut VecDeque<Message>, _context: FutureMessageContext) {
+		match message {
+			FutureMessage::Await { future } => {
+				self.spawner.spawn(future.into_future(), self.results_sender.clone(), self.wake.clone());
+			}
+			FutureMessage::Wake => {
+				// Tick-only message: the dispatcher's top-of-tick drain handles the real work.
+			}
+		}
+	}
+
+	advertise_actions!(FutureMessageDiscriminant;);
+}
+
+#[cfg(not(target_family = "wasm"))]
+fn default_spawner() -> Arc<dyn MessageSpawner> {
+	Arc::new(TokioSpawner::default())
+}
+
+#[cfg(target_family = "wasm")]
+fn default_spawner() -> Arc<dyn MessageSpawner> {
+	Arc::new(WasmSpawner)
+}
+
+#[cfg(not(target_family = "wasm"))]
+struct TokioSpawner {
+	/// Built lazily on first spawn. `multi_thread(1)` lets Tokio manage its own driver.
+	runtime: std::sync::OnceLock<tokio::runtime::Runtime>,
+}
+
+#[cfg(not(target_family = "wasm"))]
+impl Default for TokioSpawner {
+	fn default() -> Self {
+		Self { runtime: std::sync::OnceLock::new() }
+	}
+}
+
+#[cfg(not(target_family = "wasm"))]
+impl TokioSpawner {
+	fn runtime(&self) -> &tokio::runtime::Runtime {
+		self.runtime.get_or_init(|| {
+			tokio::runtime::Builder::new_multi_thread()
+				.worker_threads(1)
+				.thread_name("graphite-async")
+				.enable_all()
+				.build()
+				.expect("failed to construct async-message tokio runtime")
+		})
+	}
+}
+
+#[cfg(not(target_family = "wasm"))]
+impl MessageSpawner for TokioSpawner {
+	fn spawn(&self, future: InnerMessageFuture, results: UnboundedSender<Message>, wake: Wake) {
+		self.runtime().spawn(async move {
+			let message = future.await;
+			let _ = results.unbounded_send(message);
+			wake();
+		});
+	}
+}
+
+#[cfg(target_family = "wasm")]
+struct WasmSpawner;
+
+#[cfg(target_family = "wasm")]
+impl MessageSpawner for WasmSpawner {
+	fn spawn(&self, future: InnerMessageFuture, results: UnboundedSender<Message>, wake: Wake) {
+		wasm_bindgen_futures::spawn_local(async move {
+			let message = future.await;
+			let _ = results.unbounded_send(message);
+			wake();
+		});
+	}
+}

editor/src/messages/future/mod.rs

@@ -0,0 +1,7 @@
+mod future_message;
+mod future_message_handler;
+
+#[doc(inline)]
+pub use future_message::{FutureMessage, FutureMessageDiscriminant};
+#[doc(inline)]
+pub use future_message_handler::{FutureMessageContext, FutureMessageHandler, MessageFuture, MessageSpawner, Wake};

editor/src/messages/message.rs

@@ -24,6 +24,8 @@ pub enum Message {
 	#[child]
 	Frontend(FrontendMessage),
 	#[child]
+	Future(FutureMessage),
+	#[child]
 	InputPreprocessor(InputPreprocessorMessage),
 	#[child]
 	KeyMapping(KeyMappingMessage),

editor/src/messages/mod.rs

@@ -9,6 +9,7 @@ pub mod debug;
 pub mod defer;
 pub mod dialog;
 pub mod frontend;
+pub mod future;
 pub mod input_mapper;
 pub mod input_preprocessor;
 pub mod layout;

editor/src/messages/portfolio/document/document_message_handler.rs

@@ -933,20 +933,20 @@ impl MessageHandler<DocumentMessage, DocumentMessageContext<'_>> for DocumentMes
 				let mut document = self.clone();
 				let resources_load_handle = resource_storage.resources();
 
-				responses.add(FrontendMessage::Await {
-					future: FrontendMessageFuture::new(async move {
+				responses.add(FutureMessage::Await {
+					future: MessageFuture::new(async move {
 						document.resources.garbage_collect(document.used_resources(false).as_ref());
 						document.resources.embed_resources(resources_load_handle).await;
 
 						let content = document.serialize_document().into_bytes().into();
 
-						FrontendMessage::TriggerSaveDocument {
+						Message::Frontend(FrontendMessage::TriggerSaveDocument {
 							document_id,
 							name,
 							path,
 							folder,
 							content,
-						}
+						})
 					}),
 				});
 			}

editor/src/messages/prelude.rs

@@ -15,7 +15,8 @@ pub use crate::messages::dialog::export_dialog::{ExportDialogMessage, ExportDial
 pub use crate::messages::dialog::new_document_dialog::{NewDocumentDialogMessage, NewDocumentDialogMessageDiscriminant, NewDocumentDialogMessageHandler};
 pub use crate::messages::dialog::preferences_dialog::{PreferencesDialogMessage, PreferencesDialogMessageContext, PreferencesDialogMessageDiscriminant, PreferencesDialogMessageHandler};
 pub use crate::messages::dialog::{DialogMessage, DialogMessageContext, DialogMessageDiscriminant, DialogMessageHandler};
-pub use crate::messages::frontend::{FrontendMessage, FrontendMessageDiscriminant, FrontendMessageFuture};
+pub use crate::messages::frontend::{FrontendMessage, FrontendMessageDiscriminant};
+pub use crate::messages::future::{FutureMessage, FutureMessageContext, FutureMessageDiscriminant, FutureMessageHandler, MessageFuture, MessageSpawner, Wake};
 pub use crate::messages::input_mapper::key_mapping::{KeyMappingMessage, KeyMappingMessageContext, KeyMappingMessageDiscriminant, KeyMappingMessageHandler};
 pub use crate::messages::input_mapper::{InputMapperMessage, InputMapperMessageContext, InputMapperMessageDiscriminant, InputMapperMessageHandler};
 pub use crate::messages::input_preprocessor::{InputPreprocessorMessage, InputPreprocessorMessageContext, InputPreprocessorMessageDiscriminant, InputPreprocessorMessageHandler};