Replace async-std with tokio in PVF subsystem

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,18 @@ 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_timer::Delay;
 use parity_scale_codec::{Decode, Encode};
 use polkadot_parachain::primitives::ValidationResult;
 use std::{
-	sync::{
-		atomic::{AtomicBool, Ordering},
-		Arc,
-	},
+	path::{Path, PathBuf},
+	sync::Arc,
 	thread,
 	time::Duration,
 };
+use tokio::{io, net::UnixStream, runtime::Runtime, sync::Mutex};
 
 /// Spawns a new worker with the given program path that acts as the worker and the spawn timeout.
 ///
@@ -235,38 +228,42 @@ 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 {
+	worker_event_loop("execute", socket_path, |stream| async move {
 		let executor = Executor::new().map_err(|e| {
 			io::Error::new(io::ErrorKind::Other, format!("cannot create executor: {}", e))
 		})?;
+		let mutex = Arc::new(Mutex::new((stream, false)));
 
 		loop {
-			let (artifact_path, params, execution_timeout) = recv_request(&mut stream).await?;
+			let (artifact_path, params, execution_timeout) = {
+				let mut lock = mutex.lock().await;
+				// Unset the lock flag. We set it when either thread finishes.
+				lock.1 = false;
+				recv_request(&mut lock.0).await?
+			};
 			gum::debug!(
 				target: LOG_TARGET,
 				worker_pid = %std::process::id(),
 				"worker: validating artifact {}",
 				artifact_path.display(),
 			);
 
-			// Create a lock flag. We set it when either thread finishes.
-			let lock = 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 (mutex_2, cpu_time_start_2, execution_timeout_2) =
+				(mutex.clone(), cpu_time_start, execution_timeout);
 			let handle =
 				thread::Builder::new().name("CPU time monitor".into()).spawn(move || {
-					task::block_on(async {
+					let rt = Runtime::new().unwrap();
+					rt.block_on(async {
 						cpu_time_monitor_loop(
 							JobKind::Execute,
-							stream_2,
+							mutex_2,
 							cpu_time_start_2,
 							execution_timeout_2,
-							lock_2,
 						)
 						.await;
 					})
@@ -275,18 +272,16 @@ pub fn worker_entrypoint(socket_path: &str) {
 			let response =
 				validate_using_artifact(&artifact_path, &params, &executor, cpu_time_start).await;
 
-			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,
+			let mut lock = mutex.lock().await;
+			if lock.1 {
+				// Monitor thread detected timeout and the process should be terminated soon,
 				// nothing to do.
+				let _ = handle.join();
 				continue
 			}
+			lock.1 = true;
 
-			send_response(&mut stream, response).await?;
+			send_response(&mut lock.0, response).await?;
 		}
 	});
 }