move ext

client/executor/src/lib.rs

@@ -31,7 +31,6 @@
 #![warn(missing_docs)]
 #![recursion_limit="128"]
 
-mod async_externalities;
 #[macro_use]
 mod native_executor;
 mod wasm_runtime;

client/executor/src/native_executor.rs

@@ -19,7 +19,6 @@
 use crate::{
 	RuntimeInfo, error::{Error, Result},
 	wasm_runtime::{RuntimeCache, WasmExecutionMethod},
-	async_externalities::new_async_externalities,
 };
 use sp_version::{NativeVersion, RuntimeVersion};
 use codec::{Decode, Encode};
@@ -31,7 +30,7 @@ use std::{result, panic::{UnwindSafe, AssertUnwindSafe}, sync::Arc, collections:
 use sp_wasm_interface::{HostFunctions, Function};
 use sc_executor_common::wasm_runtime::{WasmInstance, WasmModule, CallSite};
 use sp_externalities::ExternalitiesExt as _;
-use sp_io::RuntimeSpawnExt;
+use sp_io::{RuntimeSpawnExt, new_async_externalities};
 
 /// Default num of pages for the heap
 const DEFAULT_HEAP_PAGES: u64 = 1024;
@@ -306,8 +305,10 @@ impl sp_io::RuntimeSpawn for RuntimeInstanceSpawn {
 		self.scheduler.spawn("forked runtime invoke", Box::pin(async move {
 			let module = AssertUnwindSafe(module);
 			let result = with_externalities_safe(
-				&mut new_async_externalities(scheduler, module.clone())
-					.expect("Failed to setup externalities for async context"),
+				&mut new_async_externalities(scheduler.clone())
+					.expect("Failed to setup externalities for async context")
+					.with_runtime_ext(Box::new(RuntimeInstanceSpawn::new(module.clone(), scheduler)))
+					.expect("Failed to setup runtime extension for async externalities"),
 				move || {
 
 					// FIXME: Should be refactored to shared "instance factory".

client/executor/src/async_externalities.rs → primitives/io/src/async_externalities.rs

@@ -18,45 +18,47 @@
 
 //! Async externalities.
 
-use std::{
-	any::{TypeId, Any},
-	sync::Arc,
-};
+use std::any::{TypeId, Any};
 use sp_core::{
 	storage::{ChildInfo, TrackedStorageKey},
 	traits::{Externalities, SpawnNamed, TaskExecutorExt},
 };
-use sp_io::RuntimeSpawnExt;
+use crate::{RuntimeSpawnExt, RuntimeSpawn};
 use sp_externalities::{Extensions, ExternalitiesExt as _};
-use crate::native_executor::RuntimeInstanceSpawn;
 
 /// Simple state-less externalities for use in async context.
 ///
 /// Will panic if anything is accessing the storage.
 #[derive(Debug)]
 pub struct AsyncExternalities {
-    extensions: Extensions,
+	extensions: Extensions,
 }
 
-pub fn new_async_externalities(
-	scheduler: Box<dyn SpawnNamed>,
-	module: Arc<dyn sc_executor_common::wasm_runtime::WasmModule>,
-) -> Result<AsyncExternalities, &'static str> {
+/// New Async externalities.
+pub fn new_async_externalities(scheduler: Box<dyn SpawnNamed>) -> Result<AsyncExternalities, &'static str> {
 	let extensions = Extensions::default();
 	let mut res = AsyncExternalities { extensions };
 	let mut ext = &mut res as &mut dyn Externalities;
 	ext.register_extension::<TaskExecutorExt>(TaskExecutorExt(scheduler.clone()))
 		.map_err(|_| "Failed to register task executor extension.")?;
-	ext.register_extension::<RuntimeSpawnExt>(
-		RuntimeSpawnExt(Box::new(
-			RuntimeInstanceSpawn::new(module, scheduler)
-		))
-	)
-	.map_err(|_| "Failed to register task executor extension.")?;
 
 	Ok(res)
 }
 
+impl AsyncExternalities {
+	/// Extend async externalities with the ability to spawn wasm instances.
+	pub fn with_runtime_ext(
+		mut self,
+		runtime_ext: Box<dyn RuntimeSpawn>,
+	) -> Result<Self, &'static str> {
+		let mut ext = &mut self as &mut dyn Externalities;
+		ext.register_extension::<RuntimeSpawnExt>(RuntimeSpawnExt(runtime_ext))
+			.map_err(|_| "Failed to register task executor extension.")?;
+
+		Ok(self)
+	}
+}
+
 type StorageKey = Vec<u8>;
 
 type StorageValue = Vec<u8>;

primitives/io/src/lib.rs

@@ -61,6 +61,12 @@ use sp_externalities::{ExternalitiesExt, Externalities};
 #[cfg(feature = "std")]
 mod batch_verifier;
 
+#[cfg(feature = "std")]
+mod async_externalities;
+
+#[cfg(feature = "std")]
+pub use async_externalities::{new_async_externalities, AsyncExternalities};
+
 pub mod tasks;
 
 #[cfg(feature = "std")]

primitives/io/src/tasks.rs

@@ -49,8 +49,26 @@ mod inner {
 		).ok_or("Spawn called outside of externalities context!")?;
 
 		let (sender, receiver) = mpsc::channel();
+		let extra_scheduler = scheduler.clone();
 		scheduler.spawn("parallel-runtime-spawn", Box::pin(async move {
-			let result = entry_point(data);
+			let result = match crate::new_async_externalities(extra_scheduler) {
+				Ok(mut ext) => {
+					sp_externalities::set_and_run_with_externalities(
+						&mut ext,
+						move || entry_point(data),
+					)
+				},
+				Err(e) => {
+					log::warn!(
+						target: "runtime",
+						"Unable to run async task: {}",
+						e,
+					);
+
+					return;
+				},
+			};
+
 			let _ = sender.send(result);
 		}));