fix chain subscription: less boiler plate

client/rpc/src/chain/chain_full.rs

@@ -19,10 +19,14 @@
 //! Blockchain API backend for full nodes.
 
 use super::{client_err, ChainBackend, Error};
-use crate::{chain::helpers, SubscriptionTaskExecutor};
+use crate::SubscriptionTaskExecutor;
 use std::{marker::PhantomData, sync::Arc};
 
-use futures::task::Spawn;
+use futures::{
+	stream::{self, Stream, StreamExt},
+	future,
+	task::Spawn,
+};
 use jsonrpsee::ws_server::SubscriptionSink;
 use sc_client_api::{BlockBackend, BlockchainEvents};
 use sp_blockchain::HeaderBackend;
@@ -68,27 +72,104 @@ where
 	}
 
 	fn subscribe_all_heads(&self, sink: SubscriptionSink) -> Result<(), Error> {
-		let client = self.client.clone();
-		let executor = self.executor.clone();
-
-		let fut = helpers::subscribe_headers(client, sink, "chain_subscribeAllHeads");
-		executor.spawn_obj(Box::pin(fut).into()).map_err(|e| Error::Client(Box::new(e)))
+		subscribe_headers(
+			&self.client,
+			self.executor.clone(),
+			"chain_subscribeAllHeads",
+			sink,
+			|| self.client().info().best_hash,
+			|| {
+				self.client()
+					.import_notification_stream()
+					.map(|notification| notification.header)
+			},
+		)
 	}
 
 	fn subscribe_new_heads(&self, sink: SubscriptionSink) -> Result<(), Error> {
-		let client = self.client.clone();
-		let executor = self.executor.clone();
-
-		let fut = helpers::subscribe_headers(client, sink, "chain_subscribeNewHeads");
-		executor.spawn_obj(Box::pin(fut).into()).map_err(|e| Error::Client(Box::new(e)))
+		subscribe_headers(
+			&self.client,
+			self.executor.clone(),
+			"chain_subscribeNewHeads",
+			sink,
+			|| self.client().info().best_hash,
+			|| {
+				self.client()
+					.import_notification_stream()
+					.filter(|notification| future::ready(notification.is_new_best))
+					.map(|notification| notification.header)
+			},
+		)
 	}
 
 	fn subscribe_finalized_heads(&self, sink: SubscriptionSink) -> Result<(), Error> {
-		let client = self.client.clone();
-		let executor = self.executor.clone();
+		subscribe_headers(
+			&self.client,
+			self.executor.clone(),
+			"chain_subscribeFinalizedHeads",
+			sink,
+			|| self.client().info().finalized_hash,
+			|| {
+				self.client()
+					.finality_notification_stream()
+					.map(|notification| notification.header)
+			},
+		)
+	}
+}
+
+/// Subscribe to new headers.
+fn subscribe_headers<Block, Client, F, G, S>(
+	client: &Arc<Client>,
+	executor: SubscriptionTaskExecutor,
+	method: &'static str,
+	mut sink: SubscriptionSink,
+	best_block_hash: G,
+	stream: F,
+) -> Result<(), Error>
+where
+	Block: BlockT + 'static,
+	Block::Header: Unpin,
+	Client: HeaderBackend<Block> + 'static,
+	F: FnOnce() -> S,
+	G: FnOnce() -> Block::Hash,
+	S: Stream<Item = Block::Header> + Send + 'static,
+{
+		// send current head right at the start.
+		let maybe_header = client
+			.header(BlockId::Hash(best_block_hash()))
+			.map_err(client_err)
+			.and_then(|header| {
+				header.ok_or_else(|| Error::Other("Best header missing.".to_string()))
+			})
+			.map_err(|e| {
+				log::warn!("Best header error {:?}", e);
+				e
+			})
+			.ok();
+
+		// send further subscriptions
+		let stream = stream();
+
+		// NOTE: by the time we set up the stream there might be a new best block and so there is a risk
+		// that the stream has a hole in it. The alternative would be to look up the best block *after*
+		// we set up the stream and chain it to the stream. Consuming code would need to handle
+		// duplicates at the beginning of the stream though.
+		let fut = async move {
+			stream::iter(maybe_header)
+				.chain(stream)
+				.take_while(|storage| {
+					future::ready(sink.send(&storage).map_or_else(
+						|e| {
+							log::debug!("Could not send data to subscription: {} error: {:?}", method, e);
+							false
+						},
+						|_| true,
+					))
+				})
+				.for_each(|_| future::ready(()))
+				.await;
+			};
 
-		let fut =
-			helpers::subscribe_finalized_headers(client, sink, "chain_subscribeFinalizedHeads");
 		executor.spawn_obj(Box::pin(fut).into()).map_err(|e| Error::Client(Box::new(e)))
-	}
 }

client/rpc/src/chain/mod.rs

@@ -19,7 +19,6 @@
 //! Substrate blockchain API.
 
 mod chain_full;
-mod helpers;
 
 #[cfg(test)]
 mod tests;