Simplify subsystem jobs (#2037)

* Simplify subsystem jobs

This pr simplifies the subsystem jobs interface. Instead of requiring an
extra message that is used to signal that a job should be ended, a job
now ends when the receiver returns `None`. Besides that it changes the
interface to enforce that messages to a job provide a relay parent.

* Drop ToJobTrait

* Remove FromJob

We always convert this message to FromJobCommand anyway.

polkadot/node/core/provisioner/src/lib.rs

@@ -26,99 +26,31 @@ use futures::{
 };
 use polkadot_node_subsystem::{
 	errors::{ChainApiError, RuntimeApiError},
-	messages::{
-		AllMessages, ChainApiMessage, ProvisionableData, ProvisionerInherentData,
-		ProvisionerMessage, RuntimeApiMessage,
-	},
+	messages::{ChainApiMessage, ProvisionableData, ProvisionerInherentData, ProvisionerMessage, AllMessages},
 };
 use polkadot_node_subsystem_util::{
-	self as util,
-	delegated_subsystem, FromJobCommand,
-	request_availability_cores, request_persisted_validation_data, JobTrait, ToJobTrait,
-	metrics::{self, prometheus},
+	self as util, delegated_subsystem, FromJobCommand,
+	request_availability_cores, request_persisted_validation_data, JobTrait, metrics::{self, prometheus},
 };
 use polkadot_primitives::v1::{
 	BackedCandidate, BlockNumber, CoreState, Hash, OccupiedCoreAssumption,
 	SignedAvailabilityBitfield, ValidatorIndex,
 };
-use std::{convert::TryFrom, pin::Pin};
-use std::collections::BTreeMap;
+use std::{pin::Pin, collections::BTreeMap};
 use thiserror::Error;
 
 const LOG_TARGET: &str = "provisioner";
 
 struct ProvisioningJob {
 	relay_parent: Hash,
-	sender: mpsc::Sender<FromJob>,
-	receiver: mpsc::Receiver<ToJob>,
+	sender: mpsc::Sender<FromJobCommand>,
+	receiver: mpsc::Receiver<ProvisionerMessage>,
 	provisionable_data_channels: Vec<mpsc::Sender<ProvisionableData>>,
 	backed_candidates: Vec<BackedCandidate>,
 	signed_bitfields: Vec<SignedAvailabilityBitfield>,
 	metrics: Metrics,
 }
 
-/// This enum defines the messages that the provisioner is prepared to receive.
-pub enum ToJob {
-	/// The provisioner message is the main input to the provisioner.
-	Provisioner(ProvisionerMessage),
-	/// This message indicates that the provisioner should shut itself down.
-	Stop,
-}
-
-impl ToJobTrait for ToJob {
-	const STOP: Self = Self::Stop;
-
-	fn relay_parent(&self) -> Option<Hash> {
-		match self {
-			Self::Provisioner(pm) => pm.relay_parent(),
-			Self::Stop => None,
-		}
-	}
-}
-
-impl TryFrom<AllMessages> for ToJob {
-	type Error = ();
-
-	fn try_from(msg: AllMessages) -> Result<Self, Self::Error> {
-		match msg {
-			AllMessages::Provisioner(pm) => Ok(Self::Provisioner(pm)),
-			_ => Err(()),
-		}
-	}
-}
-
-impl From<ProvisionerMessage> for ToJob {
-	fn from(pm: ProvisionerMessage) -> Self {
-		Self::Provisioner(pm)
-	}
-}
-
-enum FromJob {
-	ChainApi(ChainApiMessage),
-	Runtime(RuntimeApiMessage),
-}
-
-impl From<FromJob> for FromJobCommand {
-	fn from(from_job: FromJob) -> FromJobCommand {
-		FromJobCommand::SendMessage(match from_job {
-			FromJob::ChainApi(cam) => AllMessages::ChainApi(cam),
-			FromJob::Runtime(ram) => AllMessages::RuntimeApi(ram),
-		})
-	}
-}
-
-impl TryFrom<AllMessages> for FromJob {
-	type Error = ();
-
-	fn try_from(msg: AllMessages) -> Result<Self, Self::Error> {
-		match msg {
-			AllMessages::ChainApi(chain) => Ok(FromJob::ChainApi(chain)),
-			AllMessages::RuntimeApi(runtime) => Ok(FromJob::Runtime(runtime)),
-			_ => Err(()),
-		}
-	}
-}
-
 #[derive(Debug, Error)]
 enum Error {
 	#[error(transparent)]
@@ -141,8 +73,7 @@ enum Error {
 }
 
 impl JobTrait for ProvisioningJob {
-	type ToJob = ToJob;
-	type FromJob = FromJob;
+	type ToJob = ProvisionerMessage;
 	type Error = Error;
 	type RunArgs = ();
 	type Metrics = Metrics;
@@ -157,8 +88,8 @@ impl JobTrait for ProvisioningJob {
 		relay_parent: Hash,
 		_run_args: Self::RunArgs,
 		metrics: Self::Metrics,
-		receiver: mpsc::Receiver<ToJob>,
-		sender: mpsc::Sender<FromJob>,
+		receiver: mpsc::Receiver<ProvisionerMessage>,
+		sender: mpsc::Sender<FromJobCommand>,
 	) -> Pin<Box<dyn Future<Output = Result<(), Self::Error>> + Send>> {
 		async move {
 			let job = ProvisioningJob::new(relay_parent, metrics, sender, receiver);
@@ -175,8 +106,8 @@ impl ProvisioningJob {
 	pub fn new(
 		relay_parent: Hash,
 		metrics: Metrics,
-		sender: mpsc::Sender<FromJob>,
-		receiver: mpsc::Receiver<ToJob>,
+		sender: mpsc::Sender<FromJobCommand>,
+		receiver: mpsc::Receiver<ProvisionerMessage>,
 	) -> Self {
 		Self {
 			relay_parent,
@@ -190,13 +121,13 @@ impl ProvisioningJob {
 	}
 
 	async fn run_loop(mut self) -> Result<(), Error> {
-		while let Some(msg) = self.receiver.next().await {
-			use ProvisionerMessage::{
-				ProvisionableData, RequestBlockAuthorshipData, RequestInherentData,
-			};
+		use ProvisionerMessage::{
+			ProvisionableData, RequestBlockAuthorshipData, RequestInherentData,
+		};
 
-			match msg {
-				ToJob::Provisioner(RequestInherentData(_, return_sender)) => {
+		loop {
+			match self.receiver.next().await  {
+				Some(RequestInherentData(_, return_sender)) => {
 					let _timer = self.metrics.time_request_inherent_data();
 
 					if let Err(err) = send_inherent_data(
@@ -214,10 +145,10 @@ impl ProvisioningJob {
 						self.metrics.on_inherent_data_request(Ok(()));
 					}
 				}
-				ToJob::Provisioner(RequestBlockAuthorshipData(_, sender)) => {
+				Some(RequestBlockAuthorshipData(_, sender)) => {
 					self.provisionable_data_channels.push(sender)
 				}
-				ToJob::Provisioner(ProvisionableData(_, data)) => {
+				Some(ProvisionableData(_, data)) => {
 					let _timer = self.metrics.time_provisionable_data();
 
 					let mut bad_indices = Vec::new();
@@ -252,7 +183,7 @@ impl ProvisioningJob {
 						.map(|(_, item)| item)
 						.collect();
 				}
-				ToJob::Stop => break,
+				None => break,
 			}
 		}
 
@@ -298,7 +229,7 @@ async fn send_inherent_data(
 	bitfields: &[SignedAvailabilityBitfield],
 	candidates: &[BackedCandidate],
 	return_sender: oneshot::Sender<ProvisionerInherentData>,
-	mut from_job: mpsc::Sender<FromJob>,
+	mut from_job: mpsc::Sender<FromJobCommand>,
 ) -> Result<(), Error> {
 	let availability_cores = request_availability_cores(relay_parent, &mut from_job)
 		.await?
@@ -368,7 +299,7 @@ async fn select_candidates(
 	bitfields: &[SignedAvailabilityBitfield],
 	candidates: &[BackedCandidate],
 	relay_parent: Hash,
-	sender: &mut mpsc::Sender<FromJob>,
+	sender: &mut mpsc::Sender<FromJobCommand>,
 ) -> Result<Vec<BackedCandidate>, Error> {
 	let block_number = get_block_number_under_construction(relay_parent, sender).await?;
 
@@ -432,14 +363,14 @@ async fn select_candidates(
 #[tracing::instrument(level = "trace", skip(sender), fields(subsystem = LOG_TARGET))]
 async fn get_block_number_under_construction(
 	relay_parent: Hash,
-	sender: &mut mpsc::Sender<FromJob>,
+	sender: &mut mpsc::Sender<FromJobCommand>,
 ) -> Result<BlockNumber, Error> {
 	let (tx, rx) = oneshot::channel();
 	sender
-		.send(FromJob::ChainApi(ChainApiMessage::BlockNumber(
+		.send(AllMessages::from(ChainApiMessage::BlockNumber(
 			relay_parent,
 			tx,
-		)))
+		)).into())
 		.await
 		.map_err(|e| Error::ChainApiMessageSend(e))?;
 	match rx.await? {
@@ -558,7 +489,7 @@ impl metrics::Metrics for Metrics {
 }
 
 
-delegated_subsystem!(ProvisioningJob((), Metrics) <- ToJob as ProvisioningSubsystem);
+delegated_subsystem!(ProvisioningJob((), Metrics) <- ProvisionerMessage as ProvisioningSubsystem);
 
 #[cfg(test)]
 mod tests;

polkadot/node/core/provisioner/src/tests.rs

@@ -193,13 +193,13 @@ mod select_candidates {
 	use futures_timer::Delay;
 	use super::super::*;
 	use super::{build_occupied_core, default_bitvec, occupied_core, scheduled_core};
-	use polkadot_node_subsystem::messages::RuntimeApiRequest::{
-		AvailabilityCores, PersistedValidationData as PersistedValidationDataReq,
+	use polkadot_node_subsystem::messages::{
+		AllMessages, RuntimeApiMessage,
+		RuntimeApiRequest::{AvailabilityCores, PersistedValidationData as PersistedValidationDataReq},
 	};
 	use polkadot_primitives::v1::{
 		BlockNumber, CandidateDescriptor, CommittedCandidateReceipt, PersistedValidationData,
 	};
-	use FromJob::{ChainApi, Runtime};
 
 	const BLOCK_UNDER_PRODUCTION: BlockNumber = 128;
 
@@ -207,9 +207,9 @@ mod select_candidates {
 		overseer_factory: OverseerFactory,
 		test_factory: TestFactory,
 	) where
-		OverseerFactory: FnOnce(mpsc::Receiver<FromJob>) -> Overseer,
+		OverseerFactory: FnOnce(mpsc::Receiver<FromJobCommand>) -> Overseer,
 		Overseer: Future<Output = ()>,
-		TestFactory: FnOnce(mpsc::Sender<FromJob>) -> Test,
+		TestFactory: FnOnce(mpsc::Sender<FromJobCommand>) -> Test,
 		Test: Future<Output = ()>,
 	{
 		let (tx, rx) = mpsc::channel(64);
@@ -297,20 +297,20 @@ mod select_candidates {
 		]
 	}
 
-	async fn mock_overseer(mut receiver: mpsc::Receiver<FromJob>) {
+	async fn mock_overseer(mut receiver: mpsc::Receiver<FromJobCommand>) {
 		use ChainApiMessage::BlockNumber;
 		use RuntimeApiMessage::Request;
 
 		while let Some(from_job) = receiver.next().await {
 			match from_job {
-				ChainApi(BlockNumber(_relay_parent, tx)) => {
+				FromJobCommand::SendMessage(AllMessages::ChainApi(BlockNumber(_relay_parent, tx))) => {
 					tx.send(Ok(Some(BLOCK_UNDER_PRODUCTION - 1))).unwrap()
 				}
-				Runtime(Request(
+				FromJobCommand::SendMessage(AllMessages::RuntimeApi(Request(
 					_parent_hash,
 					PersistedValidationDataReq(_para_id, _assumption, tx),
-				)) => tx.send(Ok(Some(Default::default()))).unwrap(),
-				Runtime(Request(_parent_hash, AvailabilityCores(tx))) => {
+				))) => tx.send(Ok(Some(Default::default()))).unwrap(),
+				FromJobCommand::SendMessage(AllMessages::RuntimeApi(Request(_parent_hash, AvailabilityCores(tx)))) => {
 					tx.send(Ok(mock_availability_cores())).unwrap()
 				}
 				// non-exhaustive matches are fine for testing
@@ -321,14 +321,14 @@ mod select_candidates {
 
 	#[test]
 	fn handles_overseer_failure() {
-		let overseer = |rx: mpsc::Receiver<FromJob>| async move {
+		let overseer = |rx: mpsc::Receiver<FromJobCommand>| async move {
 			// drop the receiver so it closes and the sender can't send, then just sleep long enough that
 			// this is almost certainly not the first of the two futures to complete
 			std::mem::drop(rx);
 			Delay::new(std::time::Duration::from_secs(1)).await;
 		};
 
-		let test = |mut tx: mpsc::Sender<FromJob>| async move {
+		let test = |mut tx: mpsc::Sender<FromJobCommand>| async move {
 			// wait so that the overseer can drop the rx before we attempt to send
 			Delay::new(std::time::Duration::from_millis(50)).await;
 			let result = select_candidates(&[], &[], &[], Default::default(), &mut tx).await;
@@ -341,7 +341,7 @@ mod select_candidates {
 
 	#[test]
 	fn can_succeed() {
-		test_harness(mock_overseer, |mut tx: mpsc::Sender<FromJob>| async move {
+		test_harness(mock_overseer, |mut tx: mpsc::Sender<FromJobCommand>| async move {
 			let result = select_candidates(&[], &[], &[], Default::default(), &mut tx).await;
 			println!("{:?}", result);
 			assert!(result.is_ok());
@@ -403,7 +403,7 @@ mod select_candidates {
 			.map(|&idx| candidates[idx].clone())
 			.collect();
 
-		test_harness(mock_overseer, |mut tx: mpsc::Sender<FromJob>| async move {
+		test_harness(mock_overseer, |mut tx: mpsc::Sender<FromJobCommand>| async move {
 			let result =
 				select_candidates(&mock_cores, &[], &candidates, Default::default(), &mut tx)
 					.await;