Replace async-std with tokio in PVF subsystem

node/core/candidate-validation/src/lib.rs

@@ -624,7 +624,8 @@ trait ValidationBackend {
 			self.validate_candidate(pvf.clone(), timeout, params.encode()).await;
 
 		// If we get an AmbiguousWorkerDeath error, retry once after a brief delay, on the
-		// assumption that the conditions that caused this error may have been transient.
+		// assumption that the conditions that caused this error may have been transient. Note that
+		// this error is only a result of execution itself and not of preparation.
 		if let Err(ValidationError::InvalidCandidate(WasmInvalidCandidate::AmbiguousWorkerDeath)) =
 			validation_result
 		{

node/core/pvf/Cargo.toml

@@ -10,8 +10,6 @@ path = "bin/puppet_worker.rs"
 
 [dependencies]
 always-assert = "0.1"
-async-std = { version = "1.11.0", features = ["attributes"] }
-async-process = "1.3.0"
 assert_matches = "1.4.0"
 cpu-time = "1.0.0"
 futures = "0.3.21"
@@ -21,6 +19,7 @@ gum = { package = "tracing-gum", path = "../../gum" }
 pin-project = "1.0.9"
 rand = "0.8.5"
 tempfile = "3.3.0"
+tokio = { version = "1.22.0", features = ["fs", "process"] }
 rayon = "1.5.1"
 
 parity-scale-codec = { version = "3.1.5", default-features = false, features = ["derive"] }

node/core/pvf/src/artifacts.rs

@@ -16,10 +16,10 @@
 
 use crate::{error::PrepareError, host::PrepareResultSender};
 use always_assert::always;
-use async_std::path::{Path, PathBuf};
 use polkadot_parachain::primitives::ValidationCodeHash;
 use std::{
 	collections::HashMap,
+	path::{Path, PathBuf},
 	time::{Duration, SystemTime},
 };
 
@@ -136,8 +136,8 @@ impl Artifacts {
 	pub async fn new(cache_path: &Path) -> Self {
 		// Make sure that the cache path directory and all its parents are created.
 		// First delete the entire cache. Nodes are long-running so this should populate shortly.
-		let _ = async_std::fs::remove_dir_all(cache_path).await;
-		let _ = async_std::fs::create_dir_all(cache_path).await;
+		let _ = tokio::fs::remove_dir_all(cache_path).await;
+		let _ = tokio::fs::create_dir_all(cache_path).await;
 
 		Self { artifacts: HashMap::new() }
 	}
@@ -214,9 +214,8 @@ impl Artifacts {
 #[cfg(test)]
 mod tests {
 	use super::{ArtifactId, Artifacts};
-	use async_std::path::Path;
 	use sp_core::H256;
-	use std::str::FromStr;
+	use std::{path::Path, str::FromStr};
 
 	#[test]
 	fn from_file_name() {
@@ -252,11 +251,9 @@ mod tests {
 		);
 	}
 
-	#[test]
-	fn artifacts_removes_cache_on_startup() {
-		let fake_cache_path = async_std::task::block_on(async move {
-			crate::worker_common::tmpfile("test-cache").await.unwrap()
-		});
+	#[tokio::test]
+	async fn artifacts_removes_cache_on_startup() {
+		let fake_cache_path = crate::worker_common::tmpfile("test-cache").await.unwrap();
 		let fake_artifact_path = {
 			let mut p = fake_cache_path.clone();
 			p.push("wasmtime_0x1234567890123456789012345678901234567890123456789012345678901234");
@@ -271,7 +268,7 @@ mod tests {
 		// this should remove it and re-create.
 
 		let p = &fake_cache_path;
-		async_std::task::block_on(async { Artifacts::new(p).await });
+		Artifacts::new(p).await;
 
 		assert_eq!(std::fs::read_dir(&fake_cache_path).unwrap().count(), 0);
 

node/core/pvf/src/execute/queue.rs

@@ -24,15 +24,14 @@ use crate::{
 	worker_common::{IdleWorker, WorkerHandle},
 	InvalidCandidate, ValidationError, LOG_TARGET,
 };
-use async_std::path::PathBuf;
 use futures::{
 	channel::mpsc,
 	future::BoxFuture,
 	stream::{FuturesUnordered, StreamExt as _},
 	Future, FutureExt,
 };
 use slotmap::HopSlotMap;
-use std::{collections::VecDeque, fmt, time::Duration};
+use std::{collections::VecDeque, fmt, path::PathBuf, time::Duration};
 
 slotmap::new_key_type! { struct Worker; }
 

node/core/pvf/src/execute/worker.rs

@@ -24,25 +24,20 @@ use crate::{
 	},
 	LOG_TARGET,
 };
-use async_std::{
-	io,
-	os::unix::net::UnixStream,
-	path::{Path, PathBuf},
-	task,
-};
 use cpu_time::ProcessTime;
-use futures::FutureExt;
+use futures::{pin_mut, select_biased, FutureExt};
 use futures_timer::Delay;
 use parity_scale_codec::{Decode, Encode};
 use polkadot_parachain::primitives::ValidationResult;
 use std::{
+	path::{Path, PathBuf},
 	sync::{
 		atomic::{AtomicBool, Ordering},
 		Arc,
 	},
-	thread,
 	time::Duration,
 };
+use tokio::{io, net::UnixStream};
 
 /// Spawns a new worker with the given program path that acts as the worker and the spawn timeout.
 ///
@@ -235,10 +230,10 @@ impl Response {
 /// The entrypoint that the spawned execute worker should start with. The `socket_path` specifies
 /// the path to the socket used to communicate with the host.
 pub fn worker_entrypoint(socket_path: &str) {
-	worker_event_loop("execute", socket_path, |mut stream| async move {
-		let executor = Executor::new().map_err(|e| {
+	worker_event_loop("execute", socket_path, |rt_handle, mut stream| async move {
+		let executor = Arc::new(Executor::new().map_err(|e| {
 			io::Error::new(io::ErrorKind::Other, format!("cannot create executor: {}", e))
-		})?;
+		})?);
 
 		loop {
 			let (artifact_path, params, execution_timeout) = recv_request(&mut stream).await?;
@@ -249,52 +244,56 @@ pub fn worker_entrypoint(socket_path: &str) {
 				artifact_path.display(),
 			);
 
-			// Create a lock flag. We set it when either thread finishes.
-			let lock = Arc::new(AtomicBool::new(false));
+			// Flag used only to signal to the cpu time monitor thread that it can finish.
+			let finished_flag = Arc::new(AtomicBool::new(false));
 			let cpu_time_start = ProcessTime::now();
 
-			// Spawn a new thread that runs the CPU time monitor. Continuously wakes up from
-			// sleeping and then either sleeps for the remaining CPU time, or kills the process if
-			// we exceed the CPU timeout.
-			let (stream_2, cpu_time_start_2, execution_timeout_2, lock_2) =
-				(stream.clone(), cpu_time_start, execution_timeout, lock.clone());
-			let handle =
-				thread::Builder::new().name("CPU time monitor".into()).spawn(move || {
-					task::block_on(async {
-						cpu_time_monitor_loop(
-							JobKind::Execute,
-							stream_2,
-							cpu_time_start_2,
-							execution_timeout_2,
-							lock_2,
-						)
-						.await;
-					})
-				})?;
+			// Spawn a new thread that runs the CPU time monitor.
+			let finished_flag_2 = finished_flag.clone();
+			let thread_fut = rt_handle
+				.spawn_blocking(move || {
+					cpu_time_monitor_loop(
+						JobKind::Execute,
+						cpu_time_start,
+						execution_timeout,
+						finished_flag_2,
+					);
+				})
+				.fuse();
+			let executor_2 = executor.clone();
+			let execute_fut = rt_handle
+				.spawn_blocking(move || {
+					validate_using_artifact(&artifact_path, &params, executor_2, cpu_time_start)
+				})
+				.fuse();
 
-			let response =
-				validate_using_artifact(&artifact_path, &params, &executor, cpu_time_start).await;
+			pin_mut!(thread_fut);
+			pin_mut!(execute_fut);
 
-			let lock_result =
-				lock.compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed);
-			if lock_result.is_err() {
-				// The other thread is still sending an error response over the socket. Wait on it
-				// and return.
-				let _ = handle.join();
-				// Monitor thread detected timeout and likely already terminated the process,
-				// nothing to do.
-				continue
-			}
+			let response = select_biased! {
+				// If this future is not selected, the join handle is dropped and the thread will
+				// finish in the background.
+				join_res = thread_fut => {
+					match join_res {
+						Ok(()) => Response::TimedOut,
+						Err(e) => Response::InternalError(format!("{}", e)),
+					}
+				},
+				execute_res = execute_fut => {
+					finished_flag.store(true, Ordering::Relaxed);
+					execute_res.unwrap_or_else(|e| Response::InternalError(format!("{}", e)))
+				},
+			};
 
 			send_response(&mut stream, response).await?;
 		}
 	});
 }
 
-async fn validate_using_artifact(
+fn validate_using_artifact(
 	artifact_path: &Path,
 	params: &[u8],
-	executor: &Executor,
+	executor: Arc<Executor>,
 	cpu_time_start: ProcessTime,
 ) -> Response {
 	let descriptor_bytes = match unsafe {