Async keystore + Authority-Discovery async/await (#7000)

* Asyncify sign_with

* Asyncify generate/get keys

* Complete BareCryptoStore asyncification

* Cleanup

* Rebase

* Add Proxy

* Inject keystore proxy into extensions

* Implement some methods

* Await on send

* Cleanup

* Send result over the oneshot channel sender

* Process one future at a time

* Fix cargo stuff

* Asyncify sr25519_vrf_sign

* Cherry-pick and fix changes

* Introduce SyncCryptoStore

* SQUASH ME WITH THE first commit

* Implement into SyncCryptoStore

* Implement BareCryptoStore for KeystoreProxyAdapter

* authority-discovery

* AURA

* BABE

* finality-grandpa

* offchain-workers

* benchmarking-cli

* sp_io

* test-utils

* application-crypto

* Extensions and RPC

* Client Service

* bin

* Update cargo.lock

* Implement BareCryptoStore on proxy directly

* Simplify proxy setup

* Fix authority-discover

* Pass async keystore to authority-discovery

* Fix tests

* Use async keystore in authority-discovery

* Rename BareCryptoStore to CryptoStore

* WIP

* Remote mutable borrow in CryptoStore trait

* Implement Keystore with backends

* Remove Proxy implementation

* Fix service builder and keystore user-crates

* Fix tests

* Rework authority-discovery after refactoring

* futures::select!

* Fix multiple mut borrows in authority-discovery

* Merge fixes

* Require sync

* Restore Cargo.lock

* PR feedback - round 1

* Remove Keystore and use LocalKeystore directly

Also renamed KeystoreParams to KeystoreContainer

* Join

* Remove sync requirement

* Fix keystore tests

* Fix tests

* client/authority-discovery: Remove event stream dynamic dispatching

With authority-discovery moving from a poll based future to an `async`
future Rust has difficulties propagating the `Sync` trade through the
generated state machine.

Instead of using dynamic dispatching, use a trait parameter to specify
the DHT event stream.

* Make it compile

* Fix submit_transaction

* Fix block_on issue

* Use await in async context

* Fix manual seal keystore

* Fix authoring_blocks test

* fix aura authoring_blocks

* Try to fix tests for auth-discovery

* client/authority-discovery: Fix lookup_throttling test

* client/authority-discovery: Fix triggers_dht_get_query test

* Fix epoch_authorship_works

* client/authority-discovery: Remove timing assumption in unit test

* client/authority-discovery: Revert changes to termination test

* PR feedback

* Remove deadcode and mark test code

* Fix test_sync

* Use the correct keyring type

* Return when from_service stream is closed

* Convert SyncCryptoStore to a trait

* Fix line width

* Fix line width - take 2

* Remove unused import

* Fix keystore instantiation

* PR feedback

* Remove KeystoreContainer

* Revert "Remove KeystoreContainer"

This reverts commit ea4a37c7d74f9772b93d974e05e4498af6192730.

* Take a ref of keystore

* Move keystore to dev-dependencies

* Address some PR feedback

* Missed one

* Pass keystore reference - take 2

* client/finality-grandpa: Use `Arc<dyn CryptoStore>` instead of SyncXXX

Instead of using `SyncCryptoStorePtr` within `client/finality-grandpa`,
which is a type alias for `Arc<dyn SyncCryptoStore>`, use `Arc<dyn
CryptoStore>`. Benefits are:

1. No additional mental overhead of a `SyncCryptoStorePtr`.

2. Ability for new code to use the asynchronous methods of `CryptoStore`
instead of the synchronous `SyncCryptoStore` methods within
`client/finality-granpa` without the need for larger refactorings.

Note: This commit uses `Arc<dyn CryptoStore>` instead of
`CryptoStorePtr`, as I find the type signature more descriptive. This is
subjective and in no way required.

* Remove SyncCryptoStorePtr

* Remove KeystoreContainer & SyncCryptoStorePtr

* PR feedback

* *: Use CryptoStorePtr whereever possible

* *: Define SyncCryptoStore as a pure extension trait of CryptoStore

* Follow up to SyncCryptoStore extension trait

* Adjust docs for SyncCryptoStore as Ben suggested

* Cleanup unnecessary requirements

* sp-keystore

* Use async_std::task::block_on in keystore

* Fix block_on std requirement

* Update primitives/keystore/src/lib.rs

Co-authored-by: Max Inden <mail@max-inden.de>

* Fix wasm build

* Remove unused var

* Fix wasm compilation - take 2

* Revert async-std in keystore

* Fix indent

* Fix version and copyright

* Cleanup feature = "std"

* Auth Discovery: Ignore if from_service is cloed

* Max's suggestion

* Revert async-std usage for block_on

* Address PR feedback

* Fix example offchain worker build

* Address PR feedback

* Update Cargo.lock

* Move unused methods to test helper functions

* Restore accidentally deleted cargo.lock files

* Fix unused imports

Co-authored-by: Max Inden <mail@max-inden.de>
Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>

substrate/client/authority-discovery/src/worker.rs

@@ -19,13 +19,11 @@ use crate::{error::{Error, Result}, ServicetoWorkerMsg};
 use std::collections::{HashMap, HashSet};
 use std::convert::TryInto;
 use std::marker::PhantomData;
-use std::pin::Pin;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
 use futures::channel::mpsc;
-use futures::task::{Context, Poll};
-use futures::{Future, FutureExt, ready, Stream, StreamExt, stream::Fuse};
+use futures::{FutureExt, Stream, StreamExt, stream::Fuse};
 use futures_timer::Delay;
 
 use addr_cache::AddrCache;
@@ -47,7 +45,7 @@ use sc_network::{
 };
 use sp_authority_discovery::{AuthorityDiscoveryApi, AuthorityId, AuthoritySignature, AuthorityPair};
 use sp_core::crypto::{key_types, Pair};
-use sp_core::traits::BareCryptoStorePtr;
+use sp_keystore::CryptoStore;
 use sp_runtime::{traits::Block as BlockT, generic::BlockId};
 use sp_api::ProvideRuntimeApi;
 
@@ -77,7 +75,7 @@ const MAX_IN_FLIGHT_LOOKUPS: usize = 8;
 /// Role an authority discovery module can run as.
 pub enum Role {
 	/// Actual authority as well as a reference to its key store.
-	Authority(BareCryptoStorePtr),
+	Authority(Arc<dyn CryptoStore>),
 	/// Sentry node that guards an authority.
 	///
 	/// No reference to its key store needed, as sentry nodes don't have an identity to sign
@@ -115,7 +113,7 @@ pub enum Role {
 /// When run as a sentry node, the [`Worker`] does not publish
 /// any addresses to the DHT but still discovers validators and sentry nodes of
 /// validators, i.e. only step 2 (Discovers other authorities) is executed.
-pub struct Worker<Client, Network, Block>
+pub struct Worker<Client, Network, Block, DhtEventStream>
 where
 	Block: BlockT + 'static,
 	Network: NetworkProvider,
@@ -137,7 +135,7 @@ where
 	//   - Some(vec![a, b, c, ...]): Valid addresses were specified.
 	sentry_nodes: Option<Vec<Multiaddr>>,
 	/// Channel we receive Dht events on.
-	dht_event_rx: Pin<Box<dyn Stream<Item = DhtEvent> + Send>>,
+	dht_event_rx: DhtEventStream,
 
 	/// Interval to be proactive, publishing own addresses.
 	publish_interval: Interval,
@@ -161,14 +159,14 @@ where
 	phantom: PhantomData<Block>,
 }
 
-impl<Client, Network, Block> Worker<Client, Network, Block>
+impl<Client, Network, Block, DhtEventStream> Worker<Client, Network, Block, DhtEventStream>
 where
 	Block: BlockT + Unpin + 'static,
 	Network: NetworkProvider,
 	Client: ProvideRuntimeApi<Block> + Send + Sync + 'static + HeaderBackend<Block>,
 	<Client as ProvideRuntimeApi<Block>>::Api:
 		AuthorityDiscoveryApi<Block, Error = sp_blockchain::Error>,
-	Self: Future<Output = ()>,
+	DhtEventStream: Stream<Item = DhtEvent> + Unpin,
 {
 	/// Return a new [`Worker`].
 	///
@@ -179,7 +177,7 @@ where
 		client: Arc<Client>,
 		network: Arc<Network>,
 		sentry_nodes: Vec<MultiaddrWithPeerId>,
-		dht_event_rx: Pin<Box<dyn Stream<Item = DhtEvent> + Send>>,
+		dht_event_rx: DhtEventStream,
 		role: Role,
 		prometheus_registry: Option<prometheus_endpoint::Registry>,
 	) -> Self {
@@ -247,6 +245,72 @@ where
 		}
 	}
 
+	/// Start the worker
+	pub async fn run(mut self) {
+		loop {
+			self.start_new_lookups();
+
+			futures::select! {
+				// Process incoming events.
+				event = self.dht_event_rx.next().fuse() => {
+					if let Some(event) = event {
+						self.handle_dht_event(event).await;
+					} else {
+						// This point is reached if the network has shut down, at which point there is not
+						// much else to do than to shut down the authority discovery as well.
+						return;
+					}
+				},
+				// Handle messages from [`Service`]. Ignore if sender side is closed.
+				msg = self.from_service.select_next_some() => {
+					self.process_message_from_service(msg);
+				},
+				// Set peerset priority group to a new random set of addresses.
+				_ = self.priority_group_set_interval.next().fuse() => {
+					if let Err(e) = self.set_priority_group() {
+						error!(
+							target: LOG_TARGET,
+							"Failed to set priority group: {:?}", e,
+						);
+					}
+				},
+				// Publish own addresses.
+				_ = self.publish_interval.next().fuse() => {
+					if let Err(e) = self.publish_ext_addresses().await {
+						error!(
+							target: LOG_TARGET,
+							"Failed to publish external addresses: {:?}", e,
+						);
+					}
+				},
+				// Request addresses of authorities.
+				_ = self.query_interval.next().fuse() => {
+					if let Err(e) = self.refill_pending_lookups_queue().await {
+						error!(
+							target: LOG_TARGET,
+							"Failed to request addresses of authorities: {:?}", e,
+						);
+					}
+				},
+			}
+		}
+	}
+
+	fn process_message_from_service(&self, msg: ServicetoWorkerMsg) {
+		match msg {
+			ServicetoWorkerMsg::GetAddressesByAuthorityId(authority, sender) => {
+				let _ = sender.send(
+					self.addr_cache.get_addresses_by_authority_id(&authority).map(Clone::clone),
+				);
+			}
+			ServicetoWorkerMsg::GetAuthorityIdByPeerId(peer_id, sender) => {
+				let _ = sender.send(
+					self.addr_cache.get_authority_id_by_peer_id(&peer_id).map(Clone::clone),
+				);
+			}
+		}
+	}
+
 	fn addresses_to_publish(&self) -> impl ExactSizeIterator<Item = Multiaddr> {
 		match &self.sentry_nodes {
 			Some(addrs) => Either::Left(addrs.clone().into_iter()),
@@ -268,7 +332,7 @@ where
 	}
 
 	/// Publish either our own or if specified the public addresses of our sentry nodes.
-	fn publish_ext_addresses(&mut self) -> Result<()> {
+	async fn publish_ext_addresses(&mut self) -> Result<()> {
 		let key_store = match &self.role {
 			Role::Authority(key_store) => key_store,
 			// Only authority nodes can put addresses (their own or the ones of their sentry nodes)
@@ -291,18 +355,16 @@ where
 			.encode(&mut serialized_addresses)
 			.map_err(Error::EncodingProto)?;
 
-		let keys = Worker::get_own_public_keys_within_authority_set(
-			&key_store,
-			&self.client,
-		)?.into_iter().map(Into::into).collect::<Vec<_>>();
+		let keys = Worker::<Client, Network, Block, DhtEventStream>::get_own_public_keys_within_authority_set(
+			key_store.clone(),
+			self.client.as_ref(),
+		).await?.into_iter().map(Into::into).collect::<Vec<_>>();
 
-		let signatures = key_store.read()
-			.sign_with_all(
-				key_types::AUTHORITY_DISCOVERY,
-				keys.clone(),
-				serialized_addresses.as_slice(),
-			)
-			.map_err(|_| Error::Signing)?;
+		let signatures = key_store.sign_with_all(
+			key_types::AUTHORITY_DISCOVERY,
+			keys.clone(),
+			serialized_addresses.as_slice(),
+		).await.map_err(|_| Error::Signing)?;
 
 		for (sign_result, key) in signatures.into_iter().zip(keys) {
 			let mut signed_addresses = vec![];
@@ -327,15 +389,14 @@ where
 		Ok(())
 	}
 
-	fn refill_pending_lookups_queue(&mut self) -> Result<()> {
+	async fn refill_pending_lookups_queue(&mut self) -> Result<()> {
 		let id = BlockId::hash(self.client.info().best_hash);
 
 		let local_keys = match &self.role {
 			Role::Authority(key_store) => {
-				key_store.read()
-					.sr25519_public_keys(key_types::AUTHORITY_DISCOVERY)
-					.into_iter()
-					.collect::<HashSet<_>>()
+				key_store.sr25519_public_keys(
+					key_types::AUTHORITY_DISCOVERY
+				).await.into_iter().collect::<HashSet<_>>()
 			},
 			Role::Sentry => HashSet::new(),
 		};
@@ -387,78 +448,68 @@ where
 	}
 
 	/// Handle incoming Dht events.
-	///
-	/// Returns either:
-	///   - Poll::Pending when there are no more events to handle or
-	///   - Poll::Ready(()) when the dht event stream terminated.
-	fn handle_dht_events(&mut self, cx: &mut Context) -> Poll<()>{
-		loop {
-			match ready!(self.dht_event_rx.poll_next_unpin(cx)) {
-				Some(DhtEvent::ValueFound(v)) => {
-					if let Some(metrics) = &self.metrics {
-						metrics.dht_event_received.with_label_values(&["value_found"]).inc();
-					}
-
-					if log_enabled!(log::Level::Debug) {
-						let hashes = v.iter().map(|(hash, _value)| hash.clone());
-						debug!(
-							target: LOG_TARGET,
-							"Value for hash '{:?}' found on Dht.", hashes,
-						);
-					}
-
-					if let Err(e) = self.handle_dht_value_found_event(v) {
-						if let Some(metrics) = &self.metrics {
-							metrics.handle_value_found_event_failure.inc();
-						}
-
-						debug!(
-							target: LOG_TARGET,
-							"Failed to handle Dht value found event: {:?}", e,
-						);
-					}
+	async fn handle_dht_event(&mut self, event: DhtEvent) {
+		match event {
+			DhtEvent::ValueFound(v) => {
+				if let Some(metrics) = &self.metrics {
+					metrics.dht_event_received.with_label_values(&["value_found"]).inc();
 				}
-				Some(DhtEvent::ValueNotFound(hash)) => {
-					if let Some(metrics) = &self.metrics {
-						metrics.dht_event_received.with_label_values(&["value_not_found"]).inc();
-					}
 
-					if self.in_flight_lookups.remove(&hash).is_some() {
-						debug!(
-							target: LOG_TARGET,
-							"Value for hash '{:?}' not found on Dht.", hash
-						)
-					} else {
-						debug!(
-							target: LOG_TARGET,
-							"Received 'ValueNotFound' for unexpected hash '{:?}'.", hash
-						)
-					}
-				},
-				Some(DhtEvent::ValuePut(hash)) => {
+				if log_enabled!(log::Level::Debug) {
+					let hashes = v.iter().map(|(hash, _value)| hash.clone());
+					debug!(
+						target: LOG_TARGET,
+						"Value for hash '{:?}' found on Dht.", hashes,
+					);
+				}
+
+				if let Err(e) = self.handle_dht_value_found_event(v) {
 					if let Some(metrics) = &self.metrics {
-						metrics.dht_event_received.with_label_values(&["value_put"]).inc();
+						metrics.handle_value_found_event_failure.inc();
 					}
 
 					debug!(
 						target: LOG_TARGET,
-						"Successfully put hash '{:?}' on Dht.", hash,
-					)
-				},
-				Some(DhtEvent::ValuePutFailed(hash)) => {
-					if let Some(metrics) = &self.metrics {
-						metrics.dht_event_received.with_label_values(&["value_put_failed"]).inc();
-					}
+						"Failed to handle Dht value found event: {:?}", e,
+					);
+				}
+			}
+			DhtEvent::ValueNotFound(hash) => {
+				if let Some(metrics) = &self.metrics {
+					metrics.dht_event_received.with_label_values(&["value_not_found"]).inc();
+				}
 
+				if self.in_flight_lookups.remove(&hash).is_some() {
 					debug!(
 						target: LOG_TARGET,
-						"Failed to put hash '{:?}' on Dht.", hash
+						"Value for hash '{:?}' not found on Dht.", hash
 					)
-				},
-				None => {
-					debug!(target: LOG_TARGET, "Dht event stream terminated.");
-					return Poll::Ready(());
-				},
+				} else {
+					debug!(
+						target: LOG_TARGET,
+						"Received 'ValueNotFound' for unexpected hash '{:?}'.", hash
+					)
+				}
+			},
+			DhtEvent::ValuePut(hash) => {
+				if let Some(metrics) = &self.metrics {
+					metrics.dht_event_received.with_label_values(&["value_put"]).inc();
+				}
+
+				debug!(
+					target: LOG_TARGET,
+					"Successfully put hash '{:?}' on Dht.", hash,
+				)
+			},
+			DhtEvent::ValuePutFailed(hash) => {
+				if let Some(metrics) = &self.metrics {
+					metrics.dht_event_received.with_label_values(&["value_put_failed"]).inc();
+				}
+
+				debug!(
+					target: LOG_TARGET,
+					"Failed to put hash '{:?}' on Dht.", hash
+				)
 			}
 		}
 	}
@@ -541,7 +592,6 @@ where
 				);
 			}
 		}
-
 		Ok(())
 	}
 
@@ -551,12 +601,13 @@ where
 	// one for the upcoming session. In addition it could be participating in the current and (/ or)
 	// next authority set with two keys. The function does not return all of the local authority
 	// discovery public keys, but only the ones intersecting with the current or next authority set.
-	fn get_own_public_keys_within_authority_set(
-		key_store: &BareCryptoStorePtr,
+	async fn get_own_public_keys_within_authority_set(
+		key_store: Arc<dyn CryptoStore>,
 		client: &Client,
 	) -> Result<HashSet<AuthorityId>> {
-		let local_pub_keys = key_store.read()
+		let local_pub_keys = key_store
 			.sr25519_public_keys(key_types::AUTHORITY_DISCOVERY)
+			.await
 			.into_iter()
 			.collect::<HashSet<_>>();
 
@@ -609,86 +660,6 @@ where
 	}
 }
 
-impl<Client, Network, Block> Future for Worker<Client, Network, Block>
-where
-	Block: BlockT + Unpin + 'static,
-	Network: NetworkProvider,
-	Client: ProvideRuntimeApi<Block> + Send + Sync + 'static + HeaderBackend<Block>,
-	<Client as ProvideRuntimeApi<Block>>::Api:
-		AuthorityDiscoveryApi<Block, Error = sp_blockchain::Error>,
-{
-	type Output = ();
-
-	fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
-		// Process incoming events.
-		if let Poll::Ready(()) = self.handle_dht_events(cx) {
-			// `handle_dht_events` returns `Poll::Ready(())` when the Dht event stream terminated.
-			// Termination of the Dht event stream implies that the underlying network terminated,
-			// thus authority discovery should terminate as well.
-			return Poll::Ready(());
-		}
-
-		// Publish own addresses.
-		if let Poll::Ready(_) = self.publish_interval.poll_next_unpin(cx) {
-			// Register waker of underlying task for next interval.
-			while let Poll::Ready(_) = self.publish_interval.poll_next_unpin(cx) {}
-
-			if let Err(e) = self.publish_ext_addresses() {
-				error!(
-					target: LOG_TARGET,
-					"Failed to publish external addresses: {:?}", e,
-				);
-			}
-		}
-
-		// Request addresses of authorities, refilling the pending lookups queue.
-		if let Poll::Ready(_) = self.query_interval.poll_next_unpin(cx) {
-			// Register waker of underlying task for next interval.
-			while let Poll::Ready(_) = self.query_interval.poll_next_unpin(cx) {}
-
-			if let Err(e) = self.refill_pending_lookups_queue() {
-				error!(
-					target: LOG_TARGET,
-					"Failed to refill pending lookups queue: {:?}", e,
-				);
-			}
-		}
-
-		// Set peerset priority group to a new random set of addresses.
-		if let Poll::Ready(_) = self.priority_group_set_interval.poll_next_unpin(cx) {
-			// Register waker of underlying task for next interval.
-			while let Poll::Ready(_) = self.priority_group_set_interval.poll_next_unpin(cx) {}
-
-			if let Err(e) = self.set_priority_group() {
-				error!(
-					target: LOG_TARGET,
-					"Failed to set priority group: {:?}", e,
-				);
-			}
-		}
-
-		// Handle messages from [`Service`].
-		while let Poll::Ready(Some(msg)) = self.from_service.poll_next_unpin(cx) {
-			match msg {
-				ServicetoWorkerMsg::GetAddressesByAuthorityId(authority, sender) => {
-					let _ = sender.send(
-						self.addr_cache.get_addresses_by_authority_id(&authority).map(Clone::clone),
-					);
-				}
-				ServicetoWorkerMsg::GetAuthorityIdByPeerId(peer_id, sender) => {
-					let _ = sender.send(
-						self.addr_cache.get_authority_id_by_peer_id(&peer_id).map(Clone::clone),
-					);
-				}
-			}
-		}
-
-		self.start_new_lookups();
-
-		Poll::Pending
-	}
-}
-
 /// NetworkProvider provides [`Worker`] with all necessary hooks into the
 /// underlying Substrate networking. Using this trait abstraction instead of [`NetworkService`]
 /// directly is necessary to unit test [`Worker`].
@@ -824,7 +795,7 @@ impl Metrics {
 
 // Helper functions for unit testing.
 #[cfg(test)]
-impl<Client, Network, Block> Worker<Client, Network, Block>
+impl<Block, Client, Network, DhtEventStream> Worker<Client, Network, Block, DhtEventStream>
 where
 	Block: BlockT + 'static,
 	Network: NetworkProvider,