feat: initialize Kurdistan SDK - independent fork of Polkadot SDK

pezkuwi/node/core/pvf/src/prepare/mod.rs

@@ -0,0 +1,30 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Preparation part of pipeline
+//!
+//! The validation host spins up two processes: the queue (by running [`start_queue`]) and the pool
+//! (by running [`start_pool`]).
+//!
+//! The pool will spawn workers in new processes and those should execute pass control to
+//! `pezkuwi_node_core_pvf_worker::prepare_worker_entrypoint`.
+
+mod pool;
+mod queue;
+mod worker_interface;
+
+pub use pool::start as start_pool;
+pub use queue::{start as start_queue, FromQueue, ToQueue};

pezkuwi/node/core/pvf/src/prepare/pool.rs

@@ -0,0 +1,520 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+use super::worker_interface::{self, Outcome};
+use crate::{
+	metrics::Metrics,
+	worker_interface::{IdleWorker, WorkerHandle},
+	LOG_TARGET,
+};
+use always_assert::never;
+use futures::{
+	channel::mpsc, future::BoxFuture, stream::FuturesUnordered, Future, FutureExt, StreamExt,
+};
+use pezkuwi_node_core_pvf_common::{
+	error::{PrepareError, PrepareResult},
+	pvf::PvfPrepData,
+	SecurityStatus,
+};
+use slotmap::HopSlotMap;
+use std::{
+	fmt,
+	path::{Path, PathBuf},
+	task::Poll,
+	time::Duration,
+};
+
+slotmap::new_key_type! { pub struct Worker; }
+
+/// Messages that the pool handles.
+#[derive(Debug, PartialEq, Eq)]
+pub enum ToPool {
+	/// Request a new worker to spawn.
+	///
+	/// This request won't fail in case if the worker cannot be created. Instead, we consider
+	/// the failures transient and we try to spawn a worker after a delay.
+	///
+	/// [`FromPool::Spawned`] will be returned as soon as the worker is spawned.
+	///
+	/// The client should anticipate a [`FromPool::Rip`] message, in case the spawned worker was
+	/// stopped for some reason.
+	Spawn,
+
+	/// Kill the given worker. No-op if the given worker is not running.
+	///
+	/// [`FromPool::Rip`] won't be sent in this case. However, the client should be prepared to
+	/// receive [`FromPool::Rip`] nonetheless, since the worker may be have been ripped before
+	/// this message is processed.
+	Kill(Worker),
+
+	/// Request the given worker to start working on the given code.
+	///
+	/// Once the job either succeeded or failed, a [`FromPool::Concluded`] message will be sent
+	/// back. It's also possible that the worker dies before handling the message in which case
+	/// [`FromPool::Rip`] will be sent back.
+	///
+	/// In either case, the worker is considered busy and no further `StartWork` messages should be
+	/// sent until either `Concluded` or `Rip` message is received.
+	StartWork { worker: Worker, pvf: PvfPrepData, cache_path: PathBuf },
+}
+
+/// A message sent from pool to its client.
+#[derive(Debug)]
+pub enum FromPool {
+	/// The given worker was just spawned and is ready to be used.
+	Spawned(Worker),
+
+	/// The given worker either succeeded or failed the given job.
+	Concluded {
+		/// A key for retrieving the worker data from the pool.
+		worker: Worker,
+		/// Indicates whether the worker process was killed.
+		rip: bool,
+		/// [`Ok`] indicates that compiled artifact is successfully stored on disk.
+		/// Otherwise, an [error](PrepareError) is supplied.
+		result: PrepareResult,
+	},
+
+	/// The given worker ceased to exist.
+	Rip(Worker),
+}
+
+struct WorkerData {
+	idle: Option<IdleWorker>,
+	handle: WorkerHandle,
+}
+
+impl fmt::Debug for WorkerData {
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		write!(f, "WorkerData(pid={})", self.handle.id())
+	}
+}
+
+enum PoolEvent {
+	Spawn(IdleWorker, WorkerHandle),
+	StartWork(Worker, Outcome),
+}
+
+type Mux = FuturesUnordered<BoxFuture<'static, PoolEvent>>;
+
+struct Pool {
+	// Some variables related to the current session.
+	program_path: PathBuf,
+	cache_path: PathBuf,
+	spawn_timeout: Duration,
+	node_version: Option<String>,
+	security_status: SecurityStatus,
+
+	to_pool: mpsc::Receiver<ToPool>,
+	from_pool: mpsc::UnboundedSender<FromPool>,
+	spawned: HopSlotMap<Worker, WorkerData>,
+	mux: Mux,
+
+	metrics: Metrics,
+}
+
+/// A fatal error that warrants stopping the event loop of the pool.
+struct Fatal;
+
+async fn run(
+	Pool {
+		program_path,
+		cache_path,
+		spawn_timeout,
+		node_version,
+		security_status,
+		to_pool,
+		mut from_pool,
+		mut spawned,
+		mut mux,
+		metrics,
+	}: Pool,
+) {
+	macro_rules! break_if_fatal {
+		($expr:expr) => {
+			match $expr {
+				Err(Fatal) => break,
+				Ok(v) => v,
+			}
+		};
+	}
+
+	let mut to_pool = to_pool.fuse();
+
+	loop {
+		futures::select! {
+			to_pool = to_pool.next() => {
+				let to_pool = break_if_fatal!(to_pool.ok_or(Fatal));
+				handle_to_pool(
+					&metrics,
+					&program_path,
+					&cache_path,
+					spawn_timeout,
+					node_version.clone(),
+					security_status.clone(),
+					&mut spawned,
+					&mut mux,
+					to_pool,
+				)
+			}
+			ev = mux.select_next_some() => {
+				break_if_fatal!(handle_mux(&metrics, &mut from_pool, &mut spawned, ev))
+			}
+		}
+
+		break_if_fatal!(purge_dead(&metrics, &mut from_pool, &mut spawned).await);
+	}
+}
+
+async fn purge_dead(
+	metrics: &Metrics,
+	from_pool: &mut mpsc::UnboundedSender<FromPool>,
+	spawned: &mut HopSlotMap<Worker, WorkerData>,
+) -> Result<(), Fatal> {
+	let mut to_remove = vec![];
+	for (worker, data) in spawned.iter_mut() {
+		if data.idle.is_none() {
+			// The idle token is missing, meaning this worker is now occupied: skip it. This is
+			// because the worker process is observed by the work task and should it reach the
+			// deadline or be terminated it will be handled by the corresponding mux event.
+			continue;
+		}
+
+		if let Poll::Ready(()) = futures::poll!(&mut data.handle) {
+			// a resolved future means that the worker has terminated. Weed it out.
+			to_remove.push(worker);
+		}
+	}
+	for w in to_remove {
+		if attempt_retire(metrics, spawned, w) {
+			reply(from_pool, FromPool::Rip(w))?;
+		}
+	}
+	Ok(())
+}
+
+fn handle_to_pool(
+	metrics: &Metrics,
+	program_path: &Path,
+	cache_path: &Path,
+	spawn_timeout: Duration,
+	node_version: Option<String>,
+	security_status: SecurityStatus,
+	spawned: &mut HopSlotMap<Worker, WorkerData>,
+	mux: &mut Mux,
+	to_pool: ToPool,
+) {
+	match to_pool {
+		ToPool::Spawn => {
+			gum::debug!(target: LOG_TARGET, "spawning a new prepare worker");
+			metrics.prepare_worker().on_begin_spawn();
+			mux.push(
+				spawn_worker_task(
+					program_path.to_owned(),
+					cache_path.to_owned(),
+					spawn_timeout,
+					node_version,
+					security_status,
+				)
+				.boxed(),
+			);
+		},
+		ToPool::StartWork { worker, pvf, cache_path } => {
+			if let Some(data) = spawned.get_mut(worker) {
+				if let Some(idle) = data.idle.take() {
+					let preparation_timer = metrics.time_preparation();
+					mux.push(
+						start_work_task(
+							metrics.clone(),
+							worker,
+							idle,
+							pvf,
+							cache_path,
+							preparation_timer,
+						)
+						.boxed(),
+					);
+				} else {
+					// idle token is present after spawn and after a job is concluded;
+					// the precondition for `StartWork` is it should be sent only if all previous
+					// work items concluded;
+					// thus idle token is Some;
+					// qed.
+					never!("unexpected absence of the idle token in prepare pool");
+				}
+			} else {
+				// That's a relatively normal situation since the queue may send `start_work` and
+				// before receiving it the pool would report that the worker died.
+			}
+		},
+		ToPool::Kill(worker) => {
+			gum::debug!(target: LOG_TARGET, ?worker, "killing prepare worker");
+			// It may be absent if it were previously already removed by `purge_dead`.
+			let _ = attempt_retire(metrics, spawned, worker);
+		},
+	}
+}
+
+async fn spawn_worker_task(
+	program_path: PathBuf,
+	cache_path: PathBuf,
+	spawn_timeout: Duration,
+	node_version: Option<String>,
+	security_status: SecurityStatus,
+) -> PoolEvent {
+	use futures_timer::Delay;
+
+	loop {
+		match worker_interface::spawn(
+			&program_path,
+			&cache_path,
+			spawn_timeout,
+			node_version.as_deref(),
+			security_status.clone(),
+		)
+		.await
+		{
+			Ok((idle, handle)) => break PoolEvent::Spawn(idle, handle),
+			Err(err) => {
+				gum::warn!(target: LOG_TARGET, "failed to spawn a prepare worker: {:?}", err);
+
+				// Assume that the failure intermittent and retry after a delay.
+				Delay::new(Duration::from_secs(3)).await;
+			},
+		}
+	}
+}
+
+async fn start_work_task<Timer>(
+	metrics: Metrics,
+	worker: Worker,
+	idle: IdleWorker,
+	pvf: PvfPrepData,
+	cache_path: PathBuf,
+	_preparation_timer: Option<Timer>,
+) -> PoolEvent {
+	let outcome = worker_interface::start_work(&metrics, idle, pvf, cache_path).await;
+	PoolEvent::StartWork(worker, outcome)
+}
+
+fn handle_mux(
+	metrics: &Metrics,
+	from_pool: &mut mpsc::UnboundedSender<FromPool>,
+	spawned: &mut HopSlotMap<Worker, WorkerData>,
+	event: PoolEvent,
+) -> Result<(), Fatal> {
+	match event {
+		PoolEvent::Spawn(idle, handle) => {
+			metrics.prepare_worker().on_spawned();
+
+			let worker = spawned.insert(WorkerData { idle: Some(idle), handle });
+
+			reply(from_pool, FromPool::Spawned(worker))?;
+
+			Ok(())
+		},
+		PoolEvent::StartWork(worker, outcome) => {
+			// If we receive an outcome that the worker is unreachable or that an error occurred on
+			// the worker, we attempt to kill the worker process.
+			match outcome {
+				Outcome::Concluded { worker: idle, result } =>
+					handle_concluded_no_rip(from_pool, spawned, worker, idle, result),
+				// Return `Concluded`, but do not kill the worker since the error was on the host
+				// side.
+				Outcome::CreateTmpFileErr { worker: idle, err } => handle_concluded_no_rip(
+					from_pool,
+					spawned,
+					worker,
+					idle,
+					Err(PrepareError::CreateTmpFile(err)),
+				),
+				// Return `Concluded`, but do not kill the worker since the error was on the host
+				// side.
+				Outcome::RenameTmpFile { worker: idle, err, src, dest } => handle_concluded_no_rip(
+					from_pool,
+					spawned,
+					worker,
+					idle,
+					Err(PrepareError::RenameTmpFile { err, src, dest }),
+				),
+				// Could not clear worker cache. Kill the worker so other jobs can't see the data.
+				Outcome::ClearWorkerDir { err } => {
+					if attempt_retire(metrics, spawned, worker) {
+						reply(
+							from_pool,
+							FromPool::Concluded {
+								worker,
+								rip: true,
+								result: Err(PrepareError::ClearWorkerDir(err)),
+							},
+						)?;
+					}
+
+					Ok(())
+				},
+				Outcome::Unreachable => {
+					if attempt_retire(metrics, spawned, worker) {
+						reply(from_pool, FromPool::Rip(worker))?;
+					}
+
+					Ok(())
+				},
+				Outcome::IoErr(err) => {
+					if attempt_retire(metrics, spawned, worker) {
+						reply(
+							from_pool,
+							FromPool::Concluded {
+								worker,
+								rip: true,
+								result: Err(PrepareError::IoErr(err)),
+							},
+						)?;
+					}
+
+					Ok(())
+				},
+				// The worker might still be usable, but we kill it just in case.
+				Outcome::JobDied { err, job_pid } => {
+					if attempt_retire(metrics, spawned, worker) {
+						reply(
+							from_pool,
+							FromPool::Concluded {
+								worker,
+								rip: true,
+								result: Err(PrepareError::JobDied { err, job_pid }),
+							},
+						)?;
+					}
+
+					Ok(())
+				},
+				Outcome::TimedOut => {
+					if attempt_retire(metrics, spawned, worker) {
+						reply(
+							from_pool,
+							FromPool::Concluded {
+								worker,
+								rip: true,
+								result: Err(PrepareError::TimedOut),
+							},
+						)?;
+					}
+
+					Ok(())
+				},
+				Outcome::OutOfMemory => {
+					if attempt_retire(metrics, spawned, worker) {
+						reply(
+							from_pool,
+							FromPool::Concluded {
+								worker,
+								rip: true,
+								result: Err(PrepareError::OutOfMemory),
+							},
+						)?;
+					}
+
+					Ok(())
+				},
+			}
+		},
+	}
+}
+
+fn reply(from_pool: &mut mpsc::UnboundedSender<FromPool>, m: FromPool) -> Result<(), Fatal> {
+	from_pool.unbounded_send(m).map_err(|_| Fatal)
+}
+
+/// Removes the given worker from the registry if it there. This will lead to dropping and hence
+/// to killing the worker process.
+///
+/// Returns `true` if the worker exists and was removed and the process was killed.
+///
+/// This function takes care about counting the retired workers metric.
+fn attempt_retire(
+	metrics: &Metrics,
+	spawned: &mut HopSlotMap<Worker, WorkerData>,
+	worker: Worker,
+) -> bool {
+	if spawned.remove(worker).is_some() {
+		metrics.prepare_worker().on_retired();
+		true
+	} else {
+		false
+	}
+}
+
+/// Handles the case where we received a response. There potentially was an error, but not the fault
+/// of the worker as far as we know, so the worker should not be killed.
+///
+/// This function tries to put the idle worker back into the pool and then replies with
+/// `FromPool::Concluded` with `rip: false`.
+fn handle_concluded_no_rip(
+	from_pool: &mut mpsc::UnboundedSender<FromPool>,
+	spawned: &mut HopSlotMap<Worker, WorkerData>,
+	worker: Worker,
+	idle: IdleWorker,
+	result: PrepareResult,
+) -> Result<(), Fatal> {
+	let data = match spawned.get_mut(worker) {
+		None => {
+			// Perhaps the worker was killed meanwhile and the result is no longer relevant. We
+			// already send `Rip` when purging if we detect that the worker is dead.
+			return Ok(());
+		},
+		Some(data) => data,
+	};
+
+	// We just replace the idle worker that was loaned from this option during
+	// the work starting.
+	let old = data.idle.replace(idle);
+	never!(
+		old.is_some(),
+		"old idle worker was taken out when starting work; we only replace it here; qed"
+	);
+
+	reply(from_pool, FromPool::Concluded { worker, rip: false, result })?;
+
+	Ok(())
+}
+
+/// Spins up the pool and returns the future that should be polled to make the pool functional.
+pub fn start(
+	metrics: Metrics,
+	program_path: PathBuf,
+	cache_path: PathBuf,
+	spawn_timeout: Duration,
+	node_version: Option<String>,
+	security_status: SecurityStatus,
+) -> (mpsc::Sender<ToPool>, mpsc::UnboundedReceiver<FromPool>, impl Future<Output = ()>) {
+	let (to_pool_tx, to_pool_rx) = mpsc::channel(10);
+	let (from_pool_tx, from_pool_rx) = mpsc::unbounded();
+
+	let run = run(Pool {
+		metrics,
+		program_path,
+		cache_path,
+		spawn_timeout,
+		node_version,
+		security_status,
+		to_pool: to_pool_rx,
+		from_pool: from_pool_tx,
+		spawned: HopSlotMap::with_capacity_and_key(20),
+		mux: Mux::new(),
+	});
+
+	(to_pool_tx, from_pool_rx, run)
+}

pezkuwi/node/core/pvf/src/prepare/queue.rs

@@ -0,0 +1,796 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+//! A queue that handles requests for PVF preparation.
+
+use super::pool::{self, Worker};
+use crate::{artifacts::ArtifactId, metrics::Metrics, Priority, LOG_TARGET};
+use always_assert::{always, never};
+use futures::{channel::mpsc, stream::StreamExt as _, Future, SinkExt};
+use pezkuwi_node_core_pvf_common::{error::PrepareResult, pvf::PvfPrepData};
+use std::{
+	collections::{HashMap, VecDeque},
+	path::PathBuf,
+};
+
+#[cfg(test)]
+use std::time::Duration;
+
+/// A request to pool.
+#[derive(Debug)]
+pub enum ToQueue {
+	/// This schedules preparation of the given PVF.
+	///
+	/// Note that it is incorrect to enqueue the same PVF again without first receiving the
+	/// [`FromQueue`] response.
+	Enqueue { priority: Priority, pvf: PvfPrepData },
+}
+
+/// A response from queue.
+#[derive(Debug)]
+pub struct FromQueue {
+	/// Identifier of an artifact.
+	pub(crate) artifact_id: ArtifactId,
+	/// Outcome of the PVF processing. [`Ok`] indicates that compiled artifact
+	/// is successfully stored on disk. Otherwise, an
+	/// [error](pezkuwi_node_core_pvf_common::error::PrepareError) is supplied.
+	pub(crate) result: PrepareResult,
+}
+
+#[derive(Default)]
+struct Limits {
+	/// The maximum number of workers this pool can ever host. This is expected to be a small
+	/// number, e.g. within a dozen.
+	hard_capacity: usize,
+
+	/// The number of workers we want aim to have. If there is a critical job and we are already
+	/// at `soft_capacity`, we are allowed to grow up to `hard_capacity`. Thus this should be equal
+	/// or smaller than `hard_capacity`.
+	soft_capacity: usize,
+}
+
+impl Limits {
+	/// Returns `true` if the queue is allowed to request one more worker.
+	fn can_afford_one_more(&self, spawned_num: usize, critical: bool) -> bool {
+		let cap = if critical { self.hard_capacity } else { self.soft_capacity };
+		spawned_num < cap
+	}
+
+	/// Offer the worker back to the pool. The passed worker ID must be considered unusable unless
+	/// it wasn't taken by the pool, in which case it will be returned as `Some`.
+	fn should_cull(&mut self, spawned_num: usize) -> bool {
+		spawned_num > self.soft_capacity
+	}
+}
+
+slotmap::new_key_type! { pub struct Job; }
+
+struct JobData {
+	/// The priority of this job. Can be bumped.
+	priority: Priority,
+	pvf: PvfPrepData,
+	worker: Option<Worker>,
+}
+
+#[derive(Default)]
+struct WorkerData {
+	job: Option<Job>,
+}
+
+impl WorkerData {
+	fn is_idle(&self) -> bool {
+		self.job.is_none()
+	}
+}
+
+/// A queue structured like this is prone to starving, however, we don't care that much since we
+/// expect there is going to be a limited number of critical jobs and we don't really care if
+/// background starve.
+#[derive(Default)]
+struct Unscheduled {
+	normal: VecDeque<Job>,
+	critical: VecDeque<Job>,
+}
+
+impl Unscheduled {
+	fn queue_mut(&mut self, prio: Priority) -> &mut VecDeque<Job> {
+		match prio {
+			Priority::Normal => &mut self.normal,
+			Priority::Critical => &mut self.critical,
+		}
+	}
+
+	fn add(&mut self, prio: Priority, job: Job) {
+		self.queue_mut(prio).push_back(job);
+	}
+
+	fn readd(&mut self, prio: Priority, job: Job) {
+		self.queue_mut(prio).push_front(job);
+	}
+
+	fn is_empty(&self) -> bool {
+		self.normal.is_empty() && self.critical.is_empty()
+	}
+
+	fn next(&mut self) -> Option<Job> {
+		let mut check = |prio: Priority| self.queue_mut(prio).pop_front();
+		check(Priority::Critical).or_else(|| check(Priority::Normal))
+	}
+}
+
+struct Queue {
+	metrics: Metrics,
+
+	to_queue_rx: mpsc::Receiver<ToQueue>,
+	from_queue_tx: mpsc::UnboundedSender<FromQueue>,
+
+	to_pool_tx: mpsc::Sender<pool::ToPool>,
+	from_pool_rx: mpsc::UnboundedReceiver<pool::FromPool>,
+
+	cache_path: PathBuf,
+	limits: Limits,
+
+	jobs: slotmap::SlotMap<Job, JobData>,
+
+	/// A mapping from artifact id to a job.
+	artifact_id_to_job: HashMap<ArtifactId, Job>,
+	/// The registry of all workers.
+	workers: slotmap::SparseSecondaryMap<Worker, WorkerData>,
+	/// The number of workers requested to spawn but not yet spawned.
+	spawn_inflight: usize,
+
+	/// The jobs that are not yet scheduled. These are waiting until the next `poll` where they are
+	/// processed all at once.
+	unscheduled: Unscheduled,
+}
+
+/// A fatal error that warrants stopping the queue.
+struct Fatal;
+
+impl Queue {
+	fn new(
+		metrics: Metrics,
+		soft_capacity: usize,
+		hard_capacity: usize,
+		cache_path: PathBuf,
+		to_queue_rx: mpsc::Receiver<ToQueue>,
+		from_queue_tx: mpsc::UnboundedSender<FromQueue>,
+		to_pool_tx: mpsc::Sender<pool::ToPool>,
+		from_pool_rx: mpsc::UnboundedReceiver<pool::FromPool>,
+	) -> Self {
+		Self {
+			metrics,
+			to_queue_rx,
+			from_queue_tx,
+			to_pool_tx,
+			from_pool_rx,
+			cache_path,
+			spawn_inflight: 0,
+			limits: Limits { hard_capacity, soft_capacity },
+			jobs: slotmap::SlotMap::with_key(),
+			unscheduled: Unscheduled::default(),
+			artifact_id_to_job: HashMap::new(),
+			workers: slotmap::SparseSecondaryMap::new(),
+		}
+	}
+
+	async fn run(mut self) {
+		macro_rules! break_if_fatal {
+			($expr:expr) => {
+				if let Err(Fatal) = $expr {
+					break;
+				}
+			};
+		}
+
+		loop {
+			// biased to make it behave deterministically for tests.
+			futures::select_biased! {
+				to_queue = self.to_queue_rx.select_next_some() =>
+					break_if_fatal!(handle_to_queue(&mut self, to_queue).await),
+				from_pool = self.from_pool_rx.select_next_some() =>
+					break_if_fatal!(handle_from_pool(&mut self, from_pool).await),
+			}
+		}
+	}
+}
+
+async fn handle_to_queue(queue: &mut Queue, to_queue: ToQueue) -> Result<(), Fatal> {
+	match to_queue {
+		ToQueue::Enqueue { priority, pvf } => {
+			handle_enqueue(queue, priority, pvf).await?;
+		},
+	}
+	Ok(())
+}
+
+async fn handle_enqueue(
+	queue: &mut Queue,
+	priority: Priority,
+	pvf: PvfPrepData,
+) -> Result<(), Fatal> {
+	gum::debug!(
+		target: LOG_TARGET,
+		validation_code_hash = ?pvf.code_hash(),
+		?priority,
+		preparation_timeout = ?pvf.prep_timeout(),
+		"PVF is enqueued for preparation.",
+	);
+	queue.metrics.prepare_enqueued();
+
+	let artifact_id = ArtifactId::from_pvf_prep_data(&pvf);
+	if never!(
+		queue.artifact_id_to_job.contains_key(&artifact_id),
+		"second Enqueue sent for a known artifact"
+	) {
+		// This function is called in response to a `Enqueue` message;
+		// Precondition for `Enqueue` is that it is sent only once for a PVF;
+		// Thus this should always be `false`;
+		// qed.
+		gum::warn!(
+			target: LOG_TARGET,
+			"duplicate `enqueue` command received for {:?}",
+			artifact_id,
+		);
+		return Ok(());
+	}
+
+	let job = queue.jobs.insert(JobData { priority, pvf, worker: None });
+	queue.artifact_id_to_job.insert(artifact_id, job);
+
+	if let Some(available) = find_idle_worker(queue) {
+		// This may seem not fair (w.r.t priority) on the first glance, but it should be. This is
+		// because as soon as a worker finishes with the job it's immediately given the next one.
+		assign(queue, available, job).await?;
+	} else {
+		spawn_extra_worker(queue, priority.is_critical()).await?;
+		queue.unscheduled.add(priority, job);
+	}
+
+	Ok(())
+}
+
+fn find_idle_worker(queue: &mut Queue) -> Option<Worker> {
+	queue.workers.iter().filter(|(_, data)| data.is_idle()).map(|(k, _)| k).next()
+}
+
+async fn handle_from_pool(queue: &mut Queue, from_pool: pool::FromPool) -> Result<(), Fatal> {
+	use pool::FromPool;
+	match from_pool {
+		FromPool::Spawned(worker) => handle_worker_spawned(queue, worker).await?,
+		FromPool::Concluded { worker, rip, result } =>
+			handle_worker_concluded(queue, worker, rip, result).await?,
+		FromPool::Rip(worker) => handle_worker_rip(queue, worker).await?,
+	}
+	Ok(())
+}
+
+async fn handle_worker_spawned(queue: &mut Queue, worker: Worker) -> Result<(), Fatal> {
+	queue.workers.insert(worker, WorkerData::default());
+	queue.spawn_inflight -= 1;
+
+	if let Some(job) = queue.unscheduled.next() {
+		assign(queue, worker, job).await?;
+	}
+
+	Ok(())
+}
+
+async fn handle_worker_concluded(
+	queue: &mut Queue,
+	worker: Worker,
+	rip: bool,
+	result: PrepareResult,
+) -> Result<(), Fatal> {
+	queue.metrics.prepare_concluded();
+
+	macro_rules! never_none {
+		($expr:expr) => {
+			match $expr {
+				Some(v) => v,
+				None => {
+					// Precondition of calling this is that the `$expr` is never none;
+					// Assume the conditions holds, then this never is not hit;
+					// qed.
+					never!("never_none, {}", stringify!($expr));
+					return Ok(());
+				},
+			}
+		};
+	}
+
+	// Find out on which artifact was the worker working.
+
+	// workers are registered upon spawn and removed in one of the following cases:
+	//   1. received rip signal
+	//   2. received concluded signal with rip=true;
+	// concluded signal only comes from a spawned worker and only once;
+	// rip signal is not sent after conclusion with rip=true;
+	// the worker should be registered;
+	// this can't be None;
+	// qed.
+	let worker_data = never_none!(queue.workers.get_mut(worker));
+
+	// worker_data.job is set only by `assign` and removed only here for a worker;
+	// concluded signal only comes for a worker that was previously assigned and only once;
+	// the worker should have the job;
+	// this can't be None;
+	// qed.
+	let job = never_none!(worker_data.job.take());
+
+	// job_data is inserted upon enqueue and removed only here;
+	// as was established above, this worker was previously `assign`ed to the job;
+	// that implies that the job was enqueued;
+	// conclude signal only comes once;
+	// we are just to remove the job for the first and the only time;
+	// this can't be None;
+	// qed.
+	let job_data = never_none!(queue.jobs.remove(job));
+	let artifact_id = ArtifactId::from_pvf_prep_data(&job_data.pvf);
+
+	queue.artifact_id_to_job.remove(&artifact_id);
+
+	gum::debug!(
+		target: LOG_TARGET,
+		validation_code_hash = ?artifact_id.code_hash,
+		?worker,
+		?rip,
+		"prepare worker concluded",
+	);
+
+	reply(&mut queue.from_queue_tx, FromQueue { artifact_id, result })?;
+
+	// Figure out what to do with the worker.
+	if rip {
+		let worker_data = queue.workers.remove(worker);
+		// worker should exist, it's asserted above;
+		// qed.
+		always!(worker_data.is_some());
+
+		if !queue.unscheduled.is_empty() {
+			// That is unconditionally not critical just to not accidentally fill up
+			// the pool up to the hard cap.
+			spawn_extra_worker(queue, false).await?;
+		}
+	} else if queue.limits.should_cull(queue.workers.len() + queue.spawn_inflight) {
+		// We no longer need services of this worker. Kill it.
+		queue.workers.remove(worker);
+		send_pool(&mut queue.to_pool_tx, pool::ToPool::Kill(worker)).await?;
+	} else {
+		// see if there are more work available and schedule it.
+		if let Some(job) = queue.unscheduled.next() {
+			assign(queue, worker, job).await?;
+		}
+	}
+
+	Ok(())
+}
+
+async fn handle_worker_rip(queue: &mut Queue, worker: Worker) -> Result<(), Fatal> {
+	gum::debug!(target: LOG_TARGET, ?worker, "prepare worker ripped");
+
+	let worker_data = queue.workers.remove(worker);
+	if let Some(WorkerData { job: Some(job), .. }) = worker_data {
+		// This is an edge case where the worker ripped after we sent assignment but before it
+		// was received by the pool.
+		let priority = queue.jobs.get(job).map(|data| data.priority).unwrap_or_else(|| {
+			// job is inserted upon enqueue and removed on concluded signal;
+			// this is enclosed in the if statement that narrows the situation to before
+			// conclusion;
+			// that means that the job still exists and is known;
+			// this path cannot be hit;
+			// qed.
+			never!("the job of the ripped worker must be known but it is not");
+			Priority::Normal
+		});
+		queue.unscheduled.readd(priority, job);
+	}
+
+	// If there are still jobs left, spawn another worker to replace the ripped one (but only if it
+	// was indeed removed). That is unconditionally not critical just to not accidentally fill up
+	// the pool up to the hard cap.
+	if worker_data.is_some() && !queue.unscheduled.is_empty() {
+		spawn_extra_worker(queue, false).await?;
+	}
+	Ok(())
+}
+
+/// Spawns an extra worker if possible.
+async fn spawn_extra_worker(queue: &mut Queue, critical: bool) -> Result<(), Fatal> {
+	if queue
+		.limits
+		.can_afford_one_more(queue.workers.len() + queue.spawn_inflight, critical)
+	{
+		queue.spawn_inflight += 1;
+		send_pool(&mut queue.to_pool_tx, pool::ToPool::Spawn).await?;
+	}
+
+	Ok(())
+}
+
+/// Attaches the work to the given worker telling the poll about the job.
+async fn assign(queue: &mut Queue, worker: Worker, job: Job) -> Result<(), Fatal> {
+	let job_data = &mut queue.jobs[job];
+	job_data.worker = Some(worker);
+
+	queue.workers[worker].job = Some(job);
+
+	send_pool(
+		&mut queue.to_pool_tx,
+		pool::ToPool::StartWork {
+			worker,
+			pvf: job_data.pvf.clone(),
+			cache_path: queue.cache_path.clone(),
+		},
+	)
+	.await?;
+
+	Ok(())
+}
+
+fn reply(from_queue_tx: &mut mpsc::UnboundedSender<FromQueue>, m: FromQueue) -> Result<(), Fatal> {
+	from_queue_tx.unbounded_send(m).map_err(|_| {
+		// The host has hung up and thus it's fatal and we should shutdown ourselves.
+		Fatal
+	})
+}
+
+async fn send_pool(
+	to_pool_tx: &mut mpsc::Sender<pool::ToPool>,
+	m: pool::ToPool,
+) -> Result<(), Fatal> {
+	to_pool_tx.send(m).await.map_err(|_| {
+		// The pool has hung up and thus we are no longer are able to fulfill our duties. Shutdown.
+		Fatal
+	})
+}
+
+/// Spins up the queue and returns the future that should be polled to make the queue functional.
+pub fn start(
+	metrics: Metrics,
+	soft_capacity: usize,
+	hard_capacity: usize,
+	cache_path: PathBuf,
+	to_pool_tx: mpsc::Sender<pool::ToPool>,
+	from_pool_rx: mpsc::UnboundedReceiver<pool::FromPool>,
+) -> (mpsc::Sender<ToQueue>, mpsc::UnboundedReceiver<FromQueue>, impl Future<Output = ()>) {
+	let (to_queue_tx, to_queue_rx) = mpsc::channel(150);
+	let (from_queue_tx, from_queue_rx) = mpsc::unbounded();
+
+	let run = Queue::new(
+		metrics,
+		soft_capacity,
+		hard_capacity,
+		cache_path,
+		to_queue_rx,
+		from_queue_tx,
+		to_pool_tx,
+		from_pool_rx,
+	)
+	.run();
+
+	(to_queue_tx, from_queue_rx, run)
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::host::tests::TEST_PREPARATION_TIMEOUT;
+	use assert_matches::assert_matches;
+	use futures::{future::BoxFuture, FutureExt};
+	use pezkuwi_node_core_pvf_common::{error::PrepareError, prepare::PrepareSuccess};
+	use slotmap::SlotMap;
+	use std::task::Poll;
+
+	/// Creates a new PVF which artifact id can be uniquely identified by the given number.
+	fn pvf(discriminator: u32) -> PvfPrepData {
+		PvfPrepData::from_discriminator(discriminator)
+	}
+
+	async fn run_until<R>(
+		task: &mut (impl Future<Output = ()> + Unpin),
+		mut fut: (impl Future<Output = R> + Unpin),
+	) -> R {
+		let start = std::time::Instant::now();
+		let fut = &mut fut;
+		loop {
+			if start.elapsed() > std::time::Duration::from_secs(1) {
+				// We expect that this will take only a couple of iterations and thus to take way
+				// less than a second.
+				panic!("timeout");
+			}
+
+			if let Poll::Ready(r) = futures::poll!(&mut *fut) {
+				break r;
+			}
+
+			if futures::poll!(&mut *task).is_ready() {
+				panic!()
+			}
+		}
+	}
+
+	struct Test {
+		_tempdir: tempfile::TempDir,
+		run: BoxFuture<'static, ()>,
+		workers: SlotMap<Worker, ()>,
+		from_pool_tx: mpsc::UnboundedSender<pool::FromPool>,
+		to_pool_rx: mpsc::Receiver<pool::ToPool>,
+		to_queue_tx: mpsc::Sender<ToQueue>,
+		from_queue_rx: mpsc::UnboundedReceiver<FromQueue>,
+	}
+
+	impl Test {
+		fn new(soft_capacity: usize, hard_capacity: usize) -> Self {
+			let tempdir = tempfile::tempdir().unwrap();
+
+			let (to_pool_tx, to_pool_rx) = mpsc::channel(10);
+			let (from_pool_tx, from_pool_rx) = mpsc::unbounded();
+
+			let workers: SlotMap<Worker, ()> = SlotMap::with_key();
+
+			let (to_queue_tx, from_queue_rx, run) = start(
+				Metrics::default(),
+				soft_capacity,
+				hard_capacity,
+				tempdir.path().to_owned().into(),
+				to_pool_tx,
+				from_pool_rx,
+			);
+
+			Self {
+				_tempdir: tempdir,
+				run: run.boxed(),
+				workers,
+				from_pool_tx,
+				to_pool_rx,
+				to_queue_tx,
+				from_queue_rx,
+			}
+		}
+
+		fn send_queue(&mut self, to_queue: ToQueue) {
+			self.to_queue_tx.send(to_queue).now_or_never().unwrap().unwrap();
+		}
+
+		async fn poll_and_recv_from_queue(&mut self) -> FromQueue {
+			let from_queue_rx = &mut self.from_queue_rx;
+			run_until(&mut self.run, async { from_queue_rx.next().await.unwrap() }.boxed()).await
+		}
+
+		fn send_from_pool(&mut self, from_pool: pool::FromPool) {
+			self.from_pool_tx.send(from_pool).now_or_never().unwrap().unwrap();
+		}
+
+		async fn poll_and_recv_to_pool(&mut self) -> pool::ToPool {
+			let to_pool_rx = &mut self.to_pool_rx;
+			run_until(&mut self.run, async { to_pool_rx.next().await.unwrap() }.boxed()).await
+		}
+
+		async fn poll_ensure_to_pool_is_empty(&mut self) {
+			use futures_timer::Delay;
+
+			let to_pool_rx = &mut self.to_pool_rx;
+			run_until(
+				&mut self.run,
+				async {
+					futures::select! {
+						_ = Delay::new(Duration::from_millis(500)).fuse() => (),
+						_ = to_pool_rx.next().fuse() => {
+							panic!("to pool supposed to be empty")
+						}
+					}
+				}
+				.boxed(),
+			)
+			.await
+		}
+	}
+
+	#[tokio::test]
+	async fn properly_concludes() {
+		let mut test = Test::new(2, 2);
+
+		test.send_queue(ToQueue::Enqueue { priority: Priority::Normal, pvf: pvf(1) });
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+
+		let w = test.workers.insert(());
+		test.send_from_pool(pool::FromPool::Spawned(w));
+		test.send_from_pool(pool::FromPool::Concluded {
+			worker: w,
+			rip: false,
+			result: Ok(PrepareSuccess::default()),
+		});
+
+		assert_eq!(
+			test.poll_and_recv_from_queue().await.artifact_id,
+			ArtifactId::from_pvf_prep_data(&pvf(1))
+		);
+	}
+
+	#[tokio::test]
+	async fn dont_spawn_over_soft_limit_unless_critical() {
+		let mut test = Test::new(2, 3);
+
+		let priority = Priority::Normal;
+		test.send_queue(ToQueue::Enqueue { priority, pvf: PvfPrepData::from_discriminator(1) });
+		test.send_queue(ToQueue::Enqueue { priority, pvf: PvfPrepData::from_discriminator(2) });
+		// Start a non-precheck preparation for this one.
+		test.send_queue(ToQueue::Enqueue {
+			priority,
+			pvf: PvfPrepData::from_discriminator_and_timeout(3, TEST_PREPARATION_TIMEOUT * 3),
+		});
+
+		// Receive only two spawns.
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+
+		let w1 = test.workers.insert(());
+		let w2 = test.workers.insert(());
+
+		test.send_from_pool(pool::FromPool::Spawned(w1));
+		test.send_from_pool(pool::FromPool::Spawned(w2));
+
+		// Get two start works.
+		assert_matches!(test.poll_and_recv_to_pool().await, pool::ToPool::StartWork { .. });
+		assert_matches!(test.poll_and_recv_to_pool().await, pool::ToPool::StartWork { .. });
+
+		test.send_from_pool(pool::FromPool::Concluded {
+			worker: w1,
+			rip: false,
+			result: Ok(PrepareSuccess::default()),
+		});
+
+		assert_matches!(test.poll_and_recv_to_pool().await, pool::ToPool::StartWork { .. });
+
+		// Enqueue a critical job.
+		test.send_queue(ToQueue::Enqueue {
+			priority: Priority::Critical,
+			pvf: PvfPrepData::from_discriminator(4),
+		});
+
+		// 2 out of 2 are working, but there is a critical job incoming. That means that spawning
+		// another worker is warranted.
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+	}
+
+	#[tokio::test]
+	async fn cull_unwanted() {
+		let mut test = Test::new(1, 2);
+
+		test.send_queue(ToQueue::Enqueue {
+			priority: Priority::Normal,
+			pvf: PvfPrepData::from_discriminator(1),
+		});
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+		let w1 = test.workers.insert(());
+		test.send_from_pool(pool::FromPool::Spawned(w1));
+		assert_matches!(test.poll_and_recv_to_pool().await, pool::ToPool::StartWork { .. });
+
+		// Enqueue a critical job, which warrants spawning over the soft limit.
+		test.send_queue(ToQueue::Enqueue {
+			priority: Priority::Critical,
+			pvf: PvfPrepData::from_discriminator(2),
+		});
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+
+		// However, before the new worker had a chance to spawn, the first worker finishes with its
+		// job. The old worker will be killed while the new worker will be let live, even though
+		// it's not instantiated.
+		//
+		// That's a bit silly in this context, but in production there will be an entire pool up
+		// to the `soft_capacity` of workers and it doesn't matter which one to cull. Either way,
+		// we just check that edge case of an edge case works.
+		test.send_from_pool(pool::FromPool::Concluded {
+			worker: w1,
+			rip: false,
+			result: Ok(PrepareSuccess::default()),
+		});
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Kill(w1));
+	}
+
+	#[tokio::test]
+	async fn worker_mass_die_out_doesnt_stall_queue() {
+		let mut test = Test::new(2, 2);
+
+		let priority = Priority::Normal;
+		test.send_queue(ToQueue::Enqueue { priority, pvf: PvfPrepData::from_discriminator(1) });
+		test.send_queue(ToQueue::Enqueue { priority, pvf: PvfPrepData::from_discriminator(2) });
+		// Start a non-precheck preparation for this one.
+		test.send_queue(ToQueue::Enqueue {
+			priority,
+			pvf: PvfPrepData::from_discriminator_and_timeout(3, TEST_PREPARATION_TIMEOUT * 3),
+		});
+
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+
+		let w1 = test.workers.insert(());
+		let w2 = test.workers.insert(());
+
+		test.send_from_pool(pool::FromPool::Spawned(w1));
+		test.send_from_pool(pool::FromPool::Spawned(w2));
+
+		assert_matches!(test.poll_and_recv_to_pool().await, pool::ToPool::StartWork { .. });
+		assert_matches!(test.poll_and_recv_to_pool().await, pool::ToPool::StartWork { .. });
+
+		// Conclude worker 1 and rip it.
+		test.send_from_pool(pool::FromPool::Concluded {
+			worker: w1,
+			rip: true,
+			result: Ok(PrepareSuccess::default()),
+		});
+
+		// Since there is still work, the queue requested one extra worker to spawn to handle the
+		// remaining enqueued work items.
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+		assert_eq!(
+			test.poll_and_recv_from_queue().await.artifact_id,
+			ArtifactId::from_pvf_prep_data(&pvf(1))
+		);
+	}
+
+	#[tokio::test]
+	async fn doesnt_resurrect_ripped_worker_if_no_work() {
+		let mut test = Test::new(2, 2);
+
+		test.send_queue(ToQueue::Enqueue {
+			priority: Priority::Normal,
+			pvf: PvfPrepData::from_discriminator(1),
+		});
+
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+
+		let w1 = test.workers.insert(());
+		test.send_from_pool(pool::FromPool::Spawned(w1));
+
+		assert_matches!(test.poll_and_recv_to_pool().await, pool::ToPool::StartWork { .. });
+
+		test.send_from_pool(pool::FromPool::Concluded {
+			worker: w1,
+			rip: true,
+			result: Err(PrepareError::IoErr("test".into())),
+		});
+		test.poll_ensure_to_pool_is_empty().await;
+	}
+
+	#[tokio::test]
+	async fn rip_for_start_work() {
+		let mut test = Test::new(2, 2);
+
+		test.send_queue(ToQueue::Enqueue {
+			priority: Priority::Normal,
+			pvf: PvfPrepData::from_discriminator(1),
+		});
+
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+
+		let w1 = test.workers.insert(());
+		test.send_from_pool(pool::FromPool::Spawned(w1));
+
+		// Now, to the interesting part. After the queue normally issues the `start_work` command to
+		// the pool, before receiving the command the queue may report that the worker ripped.
+		assert_matches!(test.poll_and_recv_to_pool().await, pool::ToPool::StartWork { .. });
+		test.send_from_pool(pool::FromPool::Rip(w1));
+
+		// In this case, the pool should spawn a new worker and request it to work on the item.
+		assert_eq!(test.poll_and_recv_to_pool().await, pool::ToPool::Spawn);
+
+		let w2 = test.workers.insert(());
+		test.send_from_pool(pool::FromPool::Spawned(w2));
+		assert_matches!(test.poll_and_recv_to_pool().await, pool::ToPool::StartWork { .. });
+	}
+}

pezkuwi/node/core/pvf/src/prepare/worker_interface.rs

@@ -0,0 +1,376 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Host interface to the prepare worker.
+
+use crate::{
+	artifacts::generate_artifact_path,
+	metrics::Metrics,
+	worker_interface::{
+		clear_worker_dir_path, framed_recv, framed_send, spawn_with_program_path, IdleWorker,
+		SpawnErr, WorkerDir, WorkerHandle, JOB_TIMEOUT_WALL_CLOCK_FACTOR,
+	},
+	LOG_TARGET,
+};
+use codec::{Decode, Encode};
+use pezkuwi_node_core_pvf_common::{
+	error::{PrepareError, PrepareResult, PrepareWorkerResult},
+	prepare::{PrepareStats, PrepareSuccess, PrepareWorkerSuccess},
+	pvf::PvfPrepData,
+	worker_dir, SecurityStatus,
+};
+
+use sp_core::hexdisplay::HexDisplay;
+use std::{
+	path::{Path, PathBuf},
+	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.
+///
+/// Sends a handshake message to the worker as soon as it is spawned.
+pub async fn spawn(
+	program_path: &Path,
+	cache_path: &Path,
+	spawn_timeout: Duration,
+	node_version: Option<&str>,
+	security_status: SecurityStatus,
+) -> Result<(IdleWorker, WorkerHandle), SpawnErr> {
+	let mut extra_args = vec!["prepare-worker"];
+	if let Some(node_version) = node_version {
+		extra_args.extend_from_slice(&["--node-impl-version", node_version]);
+	}
+
+	spawn_with_program_path(
+		"prepare",
+		program_path,
+		cache_path,
+		&extra_args,
+		spawn_timeout,
+		security_status,
+	)
+	.await
+}
+
+/// Outcome of PVF preparation.
+///
+/// If the idle worker token is not returned, it means the worker must be terminated.
+pub enum Outcome {
+	/// The worker has finished the work assigned to it.
+	Concluded { worker: IdleWorker, result: PrepareResult },
+	/// The host tried to reach the worker but failed. This is most likely because the worked was
+	/// killed by the system.
+	Unreachable,
+	/// The temporary file for the artifact could not be created at the given cache path.
+	CreateTmpFileErr { worker: IdleWorker, err: String },
+	/// The response from the worker is received, but the tmp file cannot be renamed (moved) to the
+	/// final destination location.
+	RenameTmpFile {
+		worker: IdleWorker,
+		err: String,
+		// Unfortunately `PathBuf` doesn't implement `Encode`/`Decode`, so we do a fallible
+		// conversion to `Option<String>`.
+		src: Option<String>,
+		dest: Option<String>,
+	},
+	/// The worker cache could not be cleared for the given reason.
+	ClearWorkerDir { err: String },
+	/// The worker failed to finish the job until the given deadline.
+	///
+	/// The worker is no longer usable and should be killed.
+	TimedOut,
+	/// An IO error occurred while receiving the result from the worker process.
+	///
+	/// This doesn't return an idle worker instance, thus this worker is no longer usable.
+	IoErr(String),
+	/// The worker ran out of memory and is aborting. The worker should be ripped.
+	OutOfMemory,
+	/// The preparation job process died, due to OOM, a seccomp violation, or some other factor.
+	///
+	/// The worker might still be usable, but we kill it just in case.
+	JobDied { err: String, job_pid: i32 },
+}
+
+/// Given the idle token of a worker and parameters of work, communicates with the worker and
+/// returns the outcome.
+///
+/// NOTE: Returning the `TimedOut`, `IoErr` or `Unreachable` outcomes will trigger the child process
+/// being killed.
+pub async fn start_work(
+	metrics: &Metrics,
+	worker: IdleWorker,
+	pvf: PvfPrepData,
+	cache_path: PathBuf,
+) -> Outcome {
+	let IdleWorker { stream, pid, worker_dir } = worker;
+
+	gum::debug!(
+		target: LOG_TARGET,
+		worker_pid = %pid,
+		?worker_dir,
+		"starting prepare for {:?}",
+		pvf,
+	);
+
+	with_worker_dir_setup(
+		worker_dir,
+		stream,
+		pid,
+		|tmp_artifact_file, mut stream, worker_dir| async move {
+			let preparation_timeout = pvf.prep_timeout();
+
+			if let Err(err) = send_request(&mut stream, &pvf).await {
+				gum::warn!(
+					target: LOG_TARGET,
+					worker_pid = %pid,
+					"failed to send a prepare request: {:?}",
+					err,
+				);
+				return Outcome::Unreachable;
+			}
+
+			// Wait for the result from the worker, keeping in mind that there may be a timeout, the
+			// worker may get killed, or something along these lines. In that case we should
+			// propagate the error to the pool.
+			//
+			// We use a generous timeout here. This is in addition to the one in the child process,
+			// in case the child stalls. We have a wall clock timeout here in the host, but a CPU
+			// timeout in the child. We want to use CPU time because it varies less than wall clock
+			// time under load, but the CPU resources of the child can only be measured from the
+			// parent after the child process terminates.
+			let timeout = preparation_timeout * JOB_TIMEOUT_WALL_CLOCK_FACTOR;
+			let result = tokio::time::timeout(timeout, recv_response(&mut stream, pid)).await;
+
+			match result {
+				// Received bytes from worker within the time limit.
+				Ok(Ok(prepare_worker_result)) =>
+					handle_response(
+						metrics,
+						IdleWorker { stream, pid, worker_dir },
+						prepare_worker_result,
+						pid,
+						tmp_artifact_file,
+						&cache_path,
+						preparation_timeout,
+					)
+					.await,
+				Ok(Err(err)) => {
+					// Communication error within the time limit.
+					gum::warn!(
+						target: LOG_TARGET,
+						worker_pid = %pid,
+						"failed to recv a prepare response: {}",
+						err,
+					);
+					Outcome::IoErr(err.to_string())
+				},
+				Err(_) => {
+					// Timed out here on the host.
+					gum::warn!(
+						target: LOG_TARGET,
+						worker_pid = %pid,
+						"did not recv a prepare response within the time limit",
+					);
+					Outcome::TimedOut
+				},
+			}
+		},
+	)
+	.await
+}
+
+/// Handles the case where we successfully received response bytes on the host from the child.
+///
+/// Here we know the artifact exists, but is still located in a temporary file which will be cleared
+/// by [`with_worker_dir_setup`].
+async fn handle_response(
+	metrics: &Metrics,
+	worker: IdleWorker,
+	result: PrepareWorkerResult,
+	worker_pid: u32,
+	tmp_file: PathBuf,
+	cache_path: &Path,
+	preparation_timeout: Duration,
+) -> Outcome {
+	// TODO: Add `checksum` to `ArtifactPathId`. See:
+	//       https://github.com/pezkuwichain/pezkuwi-sdk/issues/122
+	let PrepareWorkerSuccess {
+		checksum,
+		stats: PrepareStats { cpu_time_elapsed, memory_stats, observed_wasm_code_len },
+	} = match result.clone() {
+		Ok(result) => result,
+		// Timed out on the child. This should already be logged by the child.
+		Err(PrepareError::TimedOut) => return Outcome::TimedOut,
+		Err(PrepareError::JobDied { err, job_pid }) => return Outcome::JobDied { err, job_pid },
+		Err(PrepareError::OutOfMemory) => return Outcome::OutOfMemory,
+		Err(err) => return Outcome::Concluded { worker, result: Err(err) },
+	};
+
+	metrics.observe_code_size(observed_wasm_code_len as usize);
+
+	if cpu_time_elapsed > preparation_timeout {
+		// The job didn't complete within the timeout.
+		gum::warn!(
+			target: LOG_TARGET,
+			%worker_pid,
+			"prepare job took {}ms cpu time, exceeded preparation timeout {}ms. Clearing WIP artifact {}",
+			cpu_time_elapsed.as_millis(),
+			preparation_timeout.as_millis(),
+			tmp_file.display(),
+		);
+		return Outcome::TimedOut;
+	}
+
+	let size = match tokio::fs::metadata(cache_path).await {
+		Ok(metadata) => metadata.len(),
+		Err(err) => {
+			gum::warn!(
+				target: LOG_TARGET,
+				?cache_path,
+				"failed to read size of the artifact: {}",
+				err,
+			);
+			return Outcome::IoErr(err.to_string());
+		},
+	};
+
+	// The file name should uniquely identify the artifact even across restarts. In case the cache
+	// for some reason is not cleared correctly, we cannot
+	// accidentally execute an artifact compiled under a different wasmtime version, host
+	// environment, etc.
+	let artifact_path = generate_artifact_path(cache_path);
+
+	gum::debug!(
+		target: LOG_TARGET,
+		%worker_pid,
+		"promoting WIP artifact {} to {}",
+		tmp_file.display(),
+		artifact_path.display(),
+	);
+
+	let outcome = match tokio::fs::rename(&tmp_file, &artifact_path).await {
+		Ok(()) => Outcome::Concluded {
+			worker,
+			result: Ok(PrepareSuccess {
+				checksum,
+				path: artifact_path,
+				size,
+				stats: PrepareStats {
+					cpu_time_elapsed,
+					memory_stats: memory_stats.clone(),
+					observed_wasm_code_len,
+				},
+			}),
+		},
+		Err(err) => {
+			gum::warn!(
+				target: LOG_TARGET,
+				%worker_pid,
+				"failed to rename the artifact from {} to {}: {:?}",
+				tmp_file.display(),
+				artifact_path.display(),
+				err,
+			);
+			Outcome::RenameTmpFile {
+				worker,
+				err: format!("{:?}", err),
+				src: tmp_file.to_str().map(String::from),
+				dest: artifact_path.to_str().map(String::from),
+			}
+		},
+	};
+
+	// If there were no errors up until now, log the memory stats for a successful preparation, if
+	// available.
+	metrics.observe_preparation_memory_metrics(memory_stats);
+
+	outcome
+}
+
+/// Create a temporary file for an artifact in the worker cache, execute the given future/closure
+/// passing the file path in, and clean up the worker cache.
+///
+/// Failure to clean up the worker cache results in an error - leaving any files here could be a
+/// security issue, and we should shut down the worker. This should be very rare.
+async fn with_worker_dir_setup<F, Fut>(
+	worker_dir: WorkerDir,
+	stream: UnixStream,
+	pid: u32,
+	f: F,
+) -> Outcome
+where
+	Fut: futures::Future<Output = Outcome>,
+	F: FnOnce(PathBuf, UnixStream, WorkerDir) -> Fut,
+{
+	// Create the tmp file here so that the child doesn't need any file creation rights. This will
+	// be cleared at the end of this function.
+	let tmp_file = worker_dir::prepare_tmp_artifact(worker_dir.path());
+	if let Err(err) = tokio::fs::File::create(&tmp_file).await {
+		gum::warn!(
+			target: LOG_TARGET,
+			worker_pid = %pid,
+			?worker_dir,
+			"failed to create a temp file for the artifact: {:?}",
+			err,
+		);
+		return Outcome::CreateTmpFileErr {
+			worker: IdleWorker { stream, pid, worker_dir },
+			err: format!("{:?}", err),
+		};
+	};
+
+	let worker_dir_path = worker_dir.path().to_owned();
+	let outcome = f(tmp_file, stream, worker_dir).await;
+
+	// Try to clear the worker dir.
+	if let Err(err) = clear_worker_dir_path(&worker_dir_path) {
+		gum::warn!(
+			target: LOG_TARGET,
+			worker_pid = %pid,
+			?worker_dir_path,
+			"failed to clear worker cache after the job: {:?}",
+			err,
+		);
+		return Outcome::ClearWorkerDir { err: format!("{:?}", err) };
+	}
+
+	outcome
+}
+
+async fn send_request(stream: &mut UnixStream, pvf: &PvfPrepData) -> io::Result<()> {
+	framed_send(stream, &pvf.encode()).await?;
+	Ok(())
+}
+
+async fn recv_response(stream: &mut UnixStream, pid: u32) -> io::Result<PrepareWorkerResult> {
+	let result = framed_recv(stream).await?;
+	let result = PrepareWorkerResult::decode(&mut &result[..]).map_err(|e| {
+		// We received invalid bytes from the worker.
+		let bound_bytes = &result[..result.len().min(4)];
+		gum::warn!(
+			target: LOG_TARGET,
+			worker_pid = %pid,
+			"received unexpected response from the prepare worker: {}",
+			HexDisplay::from(&bound_bytes),
+		);
+		io::Error::new(
+			io::ErrorKind::Other,
+			format!("prepare pvf recv_response: failed to decode result: {:?}", e),
+		)
+	})?;
+	Ok(result)
+}