Allow to specify multiple relay chain RPC urls for collator node (#1880)

* Allow specification of multiple urls for relay chain rpc nodes

* Add pooled RPC client basics

* Add list of clients to pooled client

* Improve

* Forward requests to dispatcher

* Switch clients on error

* Implement rotation logic

* Improve subscription handling

* Error handling cleanup

* Remove retry from rpc-client

* Improve naming

* Improve documentation

* Improve `ClientManager` abstraction

* Adjust zombienet test

* Add more comments

* fmt

* Apply reviewers comments

* Extract reconnection to extra method

* Add comment to reconnection method

* Clean up some dependencies

* Fix build

* fmt

* Provide alias for cli argument

* Apply review comments

* Rename P* to Relay*

* Improve zombienet test

* fmt

* Fix zombienet sleep

* Simplify zombienet test

* Reduce log clutter and fix starting position

* Do not distribute duplicated imported and finalized blocks

* fmt

* Apply code review suggestions

* Move building of relay chain interface to `cumulus-client-service`

* Refactoring to not push back into channel

* FMT

cumulus/client/relay-chain-rpc-interface/src/rpc_client.rs

@@ -14,7 +14,7 @@
 // You should have received a copy of the GNU General Public License
 // along with Cumulus.  If not, see <http://www.gnu.org/licenses/>.
 
-use backoff::{future::retry_notify, ExponentialBackoff};
+use crate::reconnecting_ws_client::ReconnectingWsClient;
 use cumulus_primitives_core::{
 	relay_chain::{
 		v2::{
@@ -23,211 +23,49 @@ use cumulus_primitives_core::{
 			PvfCheckStatement, ScrapedOnChainVotes, SessionIndex, SessionInfo, ValidationCode,
 			ValidationCodeHash, ValidatorId, ValidatorIndex, ValidatorSignature,
 		},
-		CandidateHash, Hash as PHash, Header as PHeader, InboundHrmpMessage,
+		CandidateHash, Hash as RelayHash, Header as RelayHeader, InboundHrmpMessage,
 	},
 	InboundDownwardMessage, ParaId, PersistedValidationData,
 };
 use cumulus_relay_chain_interface::{RelayChainError, RelayChainResult};
-use futures::{
-	channel::mpsc::{Receiver, Sender},
-	StreamExt,
-};
-use jsonrpsee::{
-	core::{
-		client::{Client as JsonRpcClient, ClientT, Subscription, SubscriptionClientT},
-		params::ArrayParams,
-		Error as JsonRpseeError,
-	},
-	rpc_params,
-	ws_client::WsClientBuilder,
-};
+use futures::channel::mpsc::Receiver;
+use jsonrpsee::{core::params::ArrayParams, rpc_params};
 use parity_scale_codec::{Decode, Encode};
 use polkadot_service::{BlockNumber, TaskManager};
 use sc_client_api::StorageData;
 use sc_rpc_api::{state::ReadProof, system::Health};
+use serde::de::DeserializeOwned;
 use sp_api::RuntimeVersion;
 use sp_consensus_babe::Epoch;
 use sp_core::sp_std::collections::btree_map::BTreeMap;
-use sp_runtime::DeserializeOwned;
 use sp_storage::StorageKey;
-use std::sync::Arc;
-use tokio::sync::mpsc::{
-	channel as tokio_channel, Receiver as TokioReceiver, Sender as TokioSender,
-};
-
 pub use url::Url;
 
 const LOG_TARGET: &str = "relay-chain-rpc-client";
 
-const NOTIFICATION_CHANNEL_SIZE_LIMIT: usize = 20;
-
 /// Client that maps RPC methods and deserializes results
 #[derive(Clone)]
 pub struct RelayChainRpcClient {
 	/// Websocket client to make calls
-	ws_client: Arc<JsonRpcClient>,
-
-	/// Retry strategy that should be used for requests and subscriptions
-	retry_strategy: ExponentialBackoff,
-
-	/// Channel to communicate with the RPC worker
-	to_worker_channel: TokioSender<NotificationRegisterMessage>,
-}
-
-/// Worker messages to register new notification listeners
-#[derive(Clone, Debug)]
-pub enum NotificationRegisterMessage {
-	RegisterBestHeadListener(Sender<PHeader>),
-	RegisterImportListener(Sender<PHeader>),
-	RegisterFinalizationListener(Sender<PHeader>),
-}
-
-/// Worker that should be used in combination with [`RelayChainRpcClient`]. Must be polled to distribute header notifications to listeners.
-struct RpcStreamWorker {
-	// Communication channel with the RPC client
-	client_receiver: TokioReceiver<NotificationRegisterMessage>,
-
-	// Senders to distribute incoming header notifications to
-	imported_header_listeners: Vec<Sender<PHeader>>,
-	finalized_header_listeners: Vec<Sender<PHeader>>,
-	best_header_listeners: Vec<Sender<PHeader>>,
-
-	// Incoming notification subscriptions
-	rpc_imported_header_subscription: Subscription<PHeader>,
-	rpc_finalized_header_subscription: Subscription<PHeader>,
-	rpc_best_header_subscription: Subscription<PHeader>,
+	ws_client: ReconnectingWsClient,
 }
 
 /// Entry point to create [`RelayChainRpcClient`] and start a worker that distributes notifications.
 pub async fn create_client_and_start_worker(
-	url: Url,
+	urls: Vec<Url>,
 	task_manager: &mut TaskManager,
 ) -> RelayChainResult<RelayChainRpcClient> {
-	tracing::info!(target: LOG_TARGET, url = %url.to_string(), "Initializing RPC Client");
-	let ws_client = WsClientBuilder::default().build(url.as_str()).await?;
+	let ws_client = ReconnectingWsClient::new(urls, task_manager).await?;
 
-	let best_head_stream = RelayChainRpcClient::subscribe_new_best_heads(&ws_client).await?;
-	let finalized_head_stream = RelayChainRpcClient::subscribe_finalized_heads(&ws_client).await?;
-	let imported_head_stream = RelayChainRpcClient::subscribe_imported_heads(&ws_client).await?;
-
-	let (worker, sender) =
-		RpcStreamWorker::new(imported_head_stream, best_head_stream, finalized_head_stream);
-	let client = RelayChainRpcClient::new(ws_client, sender).await?;
-
-	task_manager
-		.spawn_essential_handle()
-		.spawn("relay-chain-rpc-worker", None, worker.run());
+	let client = RelayChainRpcClient::new(ws_client).await?;
 
 	Ok(client)
 }
 
-fn handle_event_distribution(
-	event: Option<Result<PHeader, JsonRpseeError>>,
-	senders: &mut Vec<Sender<PHeader>>,
-) -> Result<(), String> {
-	match event {
-		Some(Ok(header)) => {
-			senders.retain_mut(|e| {
-				match e.try_send(header.clone()) {
-					// Receiver has been dropped, remove Sender from list.
-					Err(error) if error.is_disconnected() => false,
-					// Channel is full. This should not happen.
-					// TODO: Improve error handling here
-					// https://github.com/paritytech/cumulus/issues/1482
-					Err(error) => {
-						tracing::error!(target: LOG_TARGET, ?error, "Event distribution channel has reached its limit. This can lead to missed notifications.");
-						true
-					},
-					_ => true,
-				}
-			});
-			Ok(())
-		},
-		None => Err("RPC Subscription closed.".to_string()),
-		Some(Err(err)) => Err(format!("Error in RPC subscription: {}", err)),
-	}
-}
-
-impl RpcStreamWorker {
-	/// Create new worker. Returns the worker and a channel to register new listeners.
-	fn new(
-		import_sub: Subscription<PHeader>,
-		best_sub: Subscription<PHeader>,
-		finalized_sub: Subscription<PHeader>,
-	) -> (RpcStreamWorker, TokioSender<NotificationRegisterMessage>) {
-		let (tx, rx) = tokio_channel(100);
-		let worker = RpcStreamWorker {
-			client_receiver: rx,
-			imported_header_listeners: Vec::new(),
-			finalized_header_listeners: Vec::new(),
-			best_header_listeners: Vec::new(),
-			rpc_imported_header_subscription: import_sub,
-			rpc_best_header_subscription: best_sub,
-			rpc_finalized_header_subscription: finalized_sub,
-		};
-		(worker, tx)
-	}
-
-	/// Run this worker to drive notification streams.
-	/// The worker does two things:
-	/// 1. Listen for `NotificationRegisterMessage` and register new listeners for the notification streams
-	/// 2. Distribute incoming import, best head and finalization notifications to registered listeners.
-	///    If an error occurs during sending, the receiver has been closed and we remove the sender from the list.
-	pub async fn run(mut self) {
-		let mut import_sub = self.rpc_imported_header_subscription.fuse();
-		let mut best_head_sub = self.rpc_best_header_subscription.fuse();
-		let mut finalized_sub = self.rpc_finalized_header_subscription.fuse();
-		loop {
-			tokio::select! {
-				evt = self.client_receiver.recv() => match evt {
-					Some(NotificationRegisterMessage::RegisterBestHeadListener(tx)) => {
-						self.best_header_listeners.push(tx);
-					},
-					Some(NotificationRegisterMessage::RegisterImportListener(tx)) => {
-						self.imported_header_listeners.push(tx)
-					},
-					Some(NotificationRegisterMessage::RegisterFinalizationListener(tx)) => {
-						self.finalized_header_listeners.push(tx)
-					},
-					None => {
-						tracing::error!(target: LOG_TARGET, "RPC client receiver closed. Stopping RPC Worker.");
-						return;
-					}
-				},
-				import_event = import_sub.next() => {
-					if let Err(err) = handle_event_distribution(import_event, &mut self.imported_header_listeners) {
-						tracing::error!(target: LOG_TARGET, err, "Encountered error while processing imported header notification. Stopping RPC Worker.");
-						return;
-					}
-				},
-				best_header_event = best_head_sub.next() => {
-					if let Err(err) = handle_event_distribution(best_header_event, &mut self.best_header_listeners) {
-						tracing::error!(target: LOG_TARGET, err, "Encountered error while processing best header notification. Stopping RPC Worker.");
-						return;
-					}
-				}
-				finalized_event = finalized_sub.next() => {
-					if let Err(err) = handle_event_distribution(finalized_event, &mut self.finalized_header_listeners) {
-						tracing::error!(target: LOG_TARGET, err, "Encountered error while processing finalized header notification. Stopping RPC Worker.");
-						return;
-					}
-				}
-			}
-		}
-	}
-}
-
 impl RelayChainRpcClient {
 	/// Initialize new RPC Client.
-	async fn new(
-		ws_client: JsonRpcClient,
-		sender: TokioSender<NotificationRegisterMessage>,
-	) -> RelayChainResult<Self> {
-		let client = RelayChainRpcClient {
-			to_worker_channel: sender,
-			ws_client: Arc::new(ws_client),
-			retry_strategy: ExponentialBackoff::default(),
-		};
+	async fn new(ws_client: ReconnectingWsClient) -> RelayChainResult<Self> {
+		let client = RelayChainRpcClient { ws_client };
 
 		Ok(client)
 	}
@@ -236,7 +74,7 @@ impl RelayChainRpcClient {
 	pub async fn call_remote_runtime_function<R: Decode>(
 		&self,
 		method_name: &str,
-		hash: PHash,
+		hash: RelayHash,
 		payload: Option<impl Encode>,
 	) -> RelayChainResult<R> {
 		let payload_bytes =
@@ -286,34 +124,23 @@ impl RelayChainRpcClient {
 	) -> Result<R, RelayChainError>
 	where
 		R: DeserializeOwned + std::fmt::Debug,
-		OR: Fn(&jsonrpsee::core::Error),
+		OR: Fn(&RelayChainError),
 	{
-		retry_notify(
-			self.retry_strategy.clone(),
-			|| async {
-				self.ws_client.request(method, params.clone()).await.map_err(|err| match err {
-					JsonRpseeError::Transport(_) =>
-						backoff::Error::Transient { err, retry_after: None },
-					_ => backoff::Error::Permanent(err),
-				})
-			},
-			|error, dur| tracing::trace!(target: LOG_TARGET, %error, ?dur, "Encountered transport error, retrying."),
-		)
-		.await
-		.map_err(|err| {
+		self.ws_client.request(method, params).await.map_err(|err| {
 			trace_error(&err);
-			RelayChainError::RpcCallError(method.to_string(), err)})
+			RelayChainError::RpcCallError(method.to_string())
+		})
 	}
 
 	/// Returns information regarding the current epoch.
-	pub async fn babe_api_current_epoch(&self, at: PHash) -> Result<Epoch, RelayChainError> {
+	pub async fn babe_api_current_epoch(&self, at: RelayHash) -> Result<Epoch, RelayChainError> {
 		self.call_remote_runtime_function("BabeApi_current_epoch", at, None::<()>).await
 	}
 
 	/// Old method to fetch v1 session info.
 	pub async fn parachain_host_session_info_before_version_2(
 		&self,
-		at: PHash,
+		at: RelayHash,
 		index: SessionIndex,
 	) -> Result<Option<OldV1SessionInfo>, RelayChainError> {
 		self.call_remote_runtime_function(
@@ -327,8 +154,8 @@ impl RelayChainRpcClient {
 	/// Scrape dispute relevant from on-chain, backing votes and resolved disputes.
 	pub async fn parachain_host_on_chain_votes(
 		&self,
-		at: PHash,
-	) -> Result<Option<ScrapedOnChainVotes<PHash>>, RelayChainError> {
+		at: RelayHash,
+	) -> Result<Option<ScrapedOnChainVotes<RelayHash>>, RelayChainError> {
 		self.call_remote_runtime_function("ParachainHost_on_chain_votes", at, None::<()>)
 			.await
 	}
@@ -336,7 +163,7 @@ impl RelayChainRpcClient {
 	/// Returns code hashes of PVFs that require pre-checking by validators in the active set.
 	pub async fn parachain_host_pvfs_require_precheck(
 		&self,
-		at: PHash,
+		at: RelayHash,
 	) -> Result<Vec<ValidationCodeHash>, RelayChainError> {
 		self.call_remote_runtime_function("ParachainHost_pvfs_require_precheck", at, None::<()>)
 			.await
@@ -345,7 +172,7 @@ impl RelayChainRpcClient {
 	/// Submits a PVF pre-checking statement into the transaction pool.
 	pub async fn parachain_host_submit_pvf_check_statement(
 		&self,
-		at: PHash,
+		at: RelayHash,
 		stmt: PvfCheckStatement,
 		signature: ValidatorSignature,
 	) -> Result<(), RelayChainError> {
@@ -371,9 +198,9 @@ impl RelayChainRpcClient {
 	pub async fn state_get_read_proof(
 		&self,
 		storage_keys: Vec<StorageKey>,
-		at: Option<PHash>,
-	) -> Result<ReadProof<PHash>, RelayChainError> {
-		let params = rpc_params!(storage_keys, at);
+		at: Option<RelayHash>,
+	) -> Result<ReadProof<RelayHash>, RelayChainError> {
+		let params = rpc_params![storage_keys, at];
 		self.request("state_getReadProof", params).await
 	}
 
@@ -381,17 +208,17 @@ impl RelayChainRpcClient {
 	pub async fn state_get_storage(
 		&self,
 		storage_key: StorageKey,
-		at: Option<PHash>,
+		at: Option<RelayHash>,
 	) -> Result<Option<StorageData>, RelayChainError> {
-		let params = rpc_params!(storage_key, at);
+		let params = rpc_params![storage_key, at];
 		self.request("state_getStorage", params).await
 	}
 
 	/// Get hash of the n-th block in the canon chain.
 	///
 	/// By default returns latest block hash.
-	pub async fn chain_get_head(&self, at: Option<u64>) -> Result<PHash, RelayChainError> {
-		let params = rpc_params!(at);
+	pub async fn chain_get_head(&self, at: Option<u64>) -> Result<RelayHash, RelayChainError> {
+		let params = rpc_params![at];
 		self.request("chain_getHead", params).await
 	}
 
@@ -400,7 +227,7 @@ impl RelayChainRpcClient {
 	/// should be the successor of the number of the block.
 	pub async fn parachain_host_validator_groups(
 		&self,
-		at: PHash,
+		at: RelayHash,
 	) -> Result<(Vec<Vec<ValidatorIndex>>, GroupRotationInfo), RelayChainError> {
 		self.call_remote_runtime_function("ParachainHost_validator_groups", at, None::<()>)
 			.await
@@ -409,7 +236,7 @@ impl RelayChainRpcClient {
 	/// Get a vector of events concerning candidates that occurred within a block.
 	pub async fn parachain_host_candidate_events(
 		&self,
-		at: PHash,
+		at: RelayHash,
 	) -> Result<Vec<CandidateEvent>, RelayChainError> {
 		self.call_remote_runtime_function("ParachainHost_candidate_events", at, None::<()>)
 			.await
@@ -418,7 +245,7 @@ impl RelayChainRpcClient {
 	/// Checks if the given validation outputs pass the acceptance criteria.
 	pub async fn parachain_host_check_validation_outputs(
 		&self,
-		at: PHash,
+		at: RelayHash,
 		para_id: ParaId,
 		outputs: CandidateCommitments,
 	) -> Result<bool, RelayChainError> {
@@ -435,9 +262,9 @@ impl RelayChainRpcClient {
 	/// data hash against an expected one and yields `None` if they're not equal.
 	pub async fn parachain_host_assumed_validation_data(
 		&self,
-		at: PHash,
+		at: RelayHash,
 		para_id: ParaId,
-		expected_hash: PHash,
+		expected_hash: RelayHash,
 	) -> Result<Option<(PersistedValidationData, ValidationCodeHash)>, RelayChainError> {
 		self.call_remote_runtime_function(
 			"ParachainHost_persisted_assumed_validation_data",
@@ -448,7 +275,7 @@ impl RelayChainRpcClient {
 	}
 
 	/// Get hash of last finalized block.
-	pub async fn chain_get_finalized_head(&self) -> Result<PHash, RelayChainError> {
+	pub async fn chain_get_finalized_head(&self) -> Result<RelayHash, RelayChainError> {
 		self.request("chain_getFinalizedHead", rpc_params![]).await
 	}
 
@@ -456,8 +283,8 @@ impl RelayChainRpcClient {
 	pub async fn chain_get_block_hash(
 		&self,
 		block_number: Option<polkadot_service::BlockNumber>,
-	) -> Result<Option<PHash>, RelayChainError> {
-		let params = rpc_params!(block_number);
+	) -> Result<Option<RelayHash>, RelayChainError> {
+		let params = rpc_params![block_number];
 		self.request("chain_getBlockHash", params).await
 	}
 
@@ -468,7 +295,7 @@ impl RelayChainRpcClient {
 	/// and the para already occupies a core.
 	pub async fn parachain_host_persisted_validation_data(
 		&self,
-		at: PHash,
+		at: RelayHash,
 		para_id: ParaId,
 		occupied_core_assumption: OccupiedCoreAssumption,
 	) -> Result<Option<PersistedValidationData>, RelayChainError> {
@@ -483,7 +310,7 @@ impl RelayChainRpcClient {
 	/// Get the validation code from its hash.
 	pub async fn parachain_host_validation_code_by_hash(
 		&self,
-		at: PHash,
+		at: RelayHash,
 		validation_code_hash: ValidationCodeHash,
 	) -> Result<Option<ValidationCode>, RelayChainError> {
 		self.call_remote_runtime_function(
@@ -498,15 +325,15 @@ impl RelayChainRpcClient {
 	/// Cores are either free or occupied. Free cores can have paras assigned to them.
 	pub async fn parachain_host_availability_cores(
 		&self,
-		at: PHash,
-	) -> Result<Vec<CoreState<PHash, BlockNumber>>, RelayChainError> {
+		at: RelayHash,
+	) -> Result<Vec<CoreState<RelayHash, BlockNumber>>, RelayChainError> {
 		self.call_remote_runtime_function("ParachainHost_availability_cores", at, None::<()>)
 			.await
 	}
 
 	/// Get runtime version
-	pub async fn runtime_version(&self, at: PHash) -> Result<RuntimeVersion, RelayChainError> {
-		let params = rpc_params!(at);
+	pub async fn runtime_version(&self, at: RelayHash) -> Result<RuntimeVersion, RelayChainError> {
+		let params = rpc_params![at];
 		self.request("state_getRuntimeVersion", params).await
 	}
 
@@ -514,7 +341,7 @@ impl RelayChainRpcClient {
 	/// This is a staging method! Do not use on production runtimes!
 	pub async fn parachain_host_staging_get_disputes(
 		&self,
-		at: PHash,
+		at: RelayHash,
 	) -> Result<Vec<(SessionIndex, CandidateHash, DisputeState<BlockNumber>)>, RelayChainError> {
 		self.call_remote_runtime_function("ParachainHost_staging_get_disputes", at, None::<()>)
 			.await
@@ -522,7 +349,7 @@ impl RelayChainRpcClient {
 
 	pub async fn authority_discovery_authorities(
 		&self,
-		at: PHash,
+		at: RelayHash,
 	) -> Result<Vec<sp_authority_discovery::AuthorityId>, RelayChainError> {
 		self.call_remote_runtime_function("AuthorityDiscoveryApi_authorities", at, None::<()>)
 			.await
@@ -534,7 +361,7 @@ impl RelayChainRpcClient {
 	/// and the para already occupies a core.
 	pub async fn parachain_host_validation_code(
 		&self,
-		at: PHash,
+		at: RelayHash,
 		para_id: ParaId,
 		occupied_core_assumption: OccupiedCoreAssumption,
 	) -> Result<Option<ValidationCode>, RelayChainError> {
@@ -549,7 +376,7 @@ impl RelayChainRpcClient {
 	/// Fetch the hash of the validation code used by a para, making the given `OccupiedCoreAssumption`.
 	pub async fn parachain_host_validation_code_hash(
 		&self,
-		at: PHash,
+		at: RelayHash,
 		para_id: ParaId,
 		occupied_core_assumption: OccupiedCoreAssumption,
 	) -> Result<Option<ValidationCodeHash>, RelayChainError> {
@@ -564,7 +391,7 @@ impl RelayChainRpcClient {
 	/// Get the session info for the given session, if stored.
 	pub async fn parachain_host_session_info(
 		&self,
-		at: PHash,
+		at: RelayHash,
 		index: SessionIndex,
 	) -> Result<Option<SessionInfo>, RelayChainError> {
 		self.call_remote_runtime_function("ParachainHost_session_info", at, Some(index))
@@ -574,9 +401,9 @@ impl RelayChainRpcClient {
 	/// Get header at specified hash
 	pub async fn chain_get_header(
 		&self,
-		hash: Option<PHash>,
-	) -> Result<Option<PHeader>, RelayChainError> {
-		let params = rpc_params!(hash);
+		hash: Option<RelayHash>,
+	) -> Result<Option<RelayHeader>, RelayChainError> {
+		let params = rpc_params![hash];
 		self.request("chain_getHeader", params).await
 	}
 
@@ -584,7 +411,7 @@ impl RelayChainRpcClient {
 	/// assigned to occupied cores in `availability_cores` and `None` otherwise.
 	pub async fn parachain_host_candidate_pending_availability(
 		&self,
-		at: PHash,
+		at: RelayHash,
 		para_id: ParaId,
 	) -> Result<Option<CommittedCandidateReceipt>, RelayChainError> {
 		self.call_remote_runtime_function(
@@ -600,7 +427,7 @@ impl RelayChainRpcClient {
 	/// This can be used to instantiate a `SigningContext`.
 	pub async fn parachain_host_session_index_for_child(
 		&self,
-		at: PHash,
+		at: RelayHash,
 	) -> Result<SessionIndex, RelayChainError> {
 		self.call_remote_runtime_function("ParachainHost_session_index_for_child", at, None::<()>)
 			.await
@@ -609,7 +436,7 @@ impl RelayChainRpcClient {
 	/// Get the current validators.
 	pub async fn parachain_host_validators(
 		&self,
-		at: PHash,
+		at: RelayHash,
 	) -> Result<Vec<ValidatorId>, RelayChainError> {
 		self.call_remote_runtime_function("ParachainHost_validators", at, None::<()>)
 			.await
@@ -620,7 +447,7 @@ impl RelayChainRpcClient {
 	pub async fn parachain_host_inbound_hrmp_channels_contents(
 		&self,
 		para_id: ParaId,
-		at: PHash,
+		at: RelayHash,
 	) -> Result<BTreeMap<ParaId, Vec<InboundHrmpMessage>>, RelayChainError> {
 		self.call_remote_runtime_function(
 			"ParachainHost_inbound_hrmp_channels_contents",
@@ -634,81 +461,24 @@ impl RelayChainRpcClient {
 	pub async fn parachain_host_dmq_contents(
 		&self,
 		para_id: ParaId,
-		at: PHash,
+		at: RelayHash,
 	) -> Result<Vec<InboundDownwardMessage>, RelayChainError> {
 		self.call_remote_runtime_function("ParachainHost_dmq_contents", at, Some(para_id))
 			.await
 	}
 
 	/// Get a stream of all imported relay chain headers
-	pub async fn get_imported_heads_stream(&self) -> Result<Receiver<PHeader>, RelayChainError> {
-		let (tx, rx) = futures::channel::mpsc::channel::<PHeader>(NOTIFICATION_CHANNEL_SIZE_LIMIT);
-		self.send_register_message_to_worker(NotificationRegisterMessage::RegisterImportListener(
-			tx,
-		))?;
-		Ok(rx)
+	pub fn get_imported_heads_stream(&self) -> Result<Receiver<RelayHeader>, RelayChainError> {
+		self.ws_client.get_imported_heads_stream()
 	}
 
 	/// Get a stream of new best relay chain headers
-	pub async fn get_best_heads_stream(&self) -> Result<Receiver<PHeader>, RelayChainError> {
-		let (tx, rx) = futures::channel::mpsc::channel::<PHeader>(NOTIFICATION_CHANNEL_SIZE_LIMIT);
-		self.send_register_message_to_worker(
-			NotificationRegisterMessage::RegisterBestHeadListener(tx),
-		)?;
-		Ok(rx)
+	pub fn get_best_heads_stream(&self) -> Result<Receiver<RelayHeader>, RelayChainError> {
+		self.ws_client.get_best_heads_stream()
 	}
 
 	/// Get a stream of finalized relay chain headers
-	pub async fn get_finalized_heads_stream(&self) -> Result<Receiver<PHeader>, RelayChainError> {
-		let (tx, rx) = futures::channel::mpsc::channel::<PHeader>(NOTIFICATION_CHANNEL_SIZE_LIMIT);
-		self.send_register_message_to_worker(
-			NotificationRegisterMessage::RegisterFinalizationListener(tx),
-		)?;
-		Ok(rx)
-	}
-
-	fn send_register_message_to_worker(
-		&self,
-		message: NotificationRegisterMessage,
-	) -> Result<(), RelayChainError> {
-		self.to_worker_channel
-			.try_send(message)
-			.map_err(|e| RelayChainError::WorkerCommunicationError(e.to_string()))
-	}
-
-	async fn subscribe_imported_heads(
-		ws_client: &JsonRpcClient,
-	) -> Result<Subscription<PHeader>, RelayChainError> {
-		Ok(ws_client
-			.subscribe::<PHeader, _>(
-				"chain_subscribeAllHeads",
-				rpc_params![],
-				"chain_unsubscribeAllHeads",
-			)
-			.await?)
-	}
-
-	async fn subscribe_finalized_heads(
-		ws_client: &JsonRpcClient,
-	) -> Result<Subscription<PHeader>, RelayChainError> {
-		Ok(ws_client
-			.subscribe::<PHeader, _>(
-				"chain_subscribeFinalizedHeads",
-				rpc_params![],
-				"chain_unsubscribeFinalizedHeads",
-			)
-			.await?)
-	}
-
-	async fn subscribe_new_best_heads(
-		ws_client: &JsonRpcClient,
-	) -> Result<Subscription<PHeader>, RelayChainError> {
-		Ok(ws_client
-			.subscribe::<PHeader, _>(
-				"chain_subscribeNewHeads",
-				rpc_params![],
-				"chain_unsubscribeFinalizedHeads",
-			)
-			.await?)
+	pub fn get_finalized_heads_stream(&self) -> Result<Receiver<RelayHeader>, RelayChainError> {
+		self.ws_client.get_finalized_heads_stream()
 	}
 }