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
2026-06-13 18:41:05 +00:00 · 2022-12-15 11:42:07 +01:00
parent 50f212d8c0
commit 7cab12e9d2
22 changed files with 720 additions and 428 deletions
@@ -19,7 +19,7 @@ use core::time::Duration;
 use cumulus_primitives_core::{
 	relay_chain::{
 		v2::{CommittedCandidateReceipt, OccupiedCoreAssumption, SessionIndex, ValidatorId},
-		Hash as PHash, Header as PHeader, InboundHrmpMessage,
+		Hash as RelayHash, Header as RelayHeader, InboundHrmpMessage,
 	},
 	InboundDownwardMessage, ParaId, PersistedValidationData,
 };
@@ -34,6 +34,7 @@ use std::pin::Pin;

 pub use url::Url;

+mod reconnecting_ws_client;
 mod rpc_client;
 pub use rpc_client::{create_client_and_start_worker, RelayChainRpcClient};

@@ -58,7 +59,7 @@ impl RelayChainInterface for RelayChainRpcInterface {
 	async fn retrieve_dmq_contents(
 		&self,
 		para_id: ParaId,
-		relay_parent: PHash,
+		relay_parent: RelayHash,
 	) -> RelayChainResult<Vec<InboundDownwardMessage>> {
 		self.rpc_client.parachain_host_dmq_contents(para_id, relay_parent).await
 	}
@@ -66,7 +67,7 @@ impl RelayChainInterface for RelayChainRpcInterface {
 	async fn retrieve_all_inbound_hrmp_channel_contents(
 		&self,
 		para_id: ParaId,
-		relay_parent: PHash,
+		relay_parent: RelayHash,
 	) -> RelayChainResult<BTreeMap<ParaId, Vec<InboundHrmpMessage>>> {
 		self.rpc_client
 			.parachain_host_inbound_hrmp_channels_contents(para_id, relay_parent)
@@ -75,7 +76,7 @@ impl RelayChainInterface for RelayChainRpcInterface {

 	async fn persisted_validation_data(
 		&self,
-		hash: PHash,
+		hash: RelayHash,
 		para_id: ParaId,
 		occupied_core_assumption: OccupiedCoreAssumption,
 	) -> RelayChainResult<Option<PersistedValidationData>> {
@@ -86,7 +87,7 @@ impl RelayChainInterface for RelayChainRpcInterface {

 	async fn candidate_pending_availability(
 		&self,
-		hash: PHash,
+		hash: RelayHash,
 		para_id: ParaId,
 	) -> RelayChainResult<Option<CommittedCandidateReceipt>> {
 		self.rpc_client
@@ -94,31 +95,31 @@ impl RelayChainInterface for RelayChainRpcInterface {
 			.await
 	}

-	async fn session_index_for_child(&self, hash: PHash) -> RelayChainResult<SessionIndex> {
+	async fn session_index_for_child(&self, hash: RelayHash) -> RelayChainResult<SessionIndex> {
 		self.rpc_client.parachain_host_session_index_for_child(hash).await
 	}

-	async fn validators(&self, block_id: PHash) -> RelayChainResult<Vec<ValidatorId>> {
+	async fn validators(&self, block_id: RelayHash) -> RelayChainResult<Vec<ValidatorId>> {
 		self.rpc_client.parachain_host_validators(block_id).await
 	}

 	async fn import_notification_stream(
 		&self,
-	) -> RelayChainResult<Pin<Box<dyn Stream<Item = PHeader> + Send>>> {
-		let imported_headers_stream = self.rpc_client.get_imported_heads_stream().await?;
+	) -> RelayChainResult<Pin<Box<dyn Stream<Item = RelayHeader> + Send>>> {
+		let imported_headers_stream = self.rpc_client.get_imported_heads_stream()?;

 		Ok(imported_headers_stream.boxed())
 	}

 	async fn finality_notification_stream(
 		&self,
-	) -> RelayChainResult<Pin<Box<dyn Stream<Item = PHeader> + Send>>> {
-		let imported_headers_stream = self.rpc_client.get_finalized_heads_stream().await?;
+	) -> RelayChainResult<Pin<Box<dyn Stream<Item = RelayHeader> + Send>>> {
+		let imported_headers_stream = self.rpc_client.get_finalized_heads_stream()?;

 		Ok(imported_headers_stream.boxed())
 	}

-	async fn best_block_hash(&self) -> RelayChainResult<PHash> {
+	async fn best_block_hash(&self) -> RelayChainResult<RelayHash> {
 		self.rpc_client.chain_get_head(None).await
 	}

@@ -132,7 +133,7 @@ impl RelayChainInterface for RelayChainRpcInterface {

 	async fn get_storage_by_key(
 		&self,
-		relay_parent: PHash,
+		relay_parent: RelayHash,
 		key: &[u8],
 	) -> RelayChainResult<Option<StorageValue>> {
 		let storage_key = StorageKey(key.to_vec());
@@ -144,7 +145,7 @@ impl RelayChainInterface for RelayChainRpcInterface {

 	async fn prove_read(
 		&self,
-		relay_parent: PHash,
+		relay_parent: RelayHash,
 		relevant_keys: &Vec<Vec<u8>>,
 	) -> RelayChainResult<StorageProof> {
 		let cloned = relevant_keys.clone();
@@ -167,8 +168,8 @@ impl RelayChainInterface for RelayChainRpcInterface {
 	/// 2. Check if the block is already in chain. If yes, succeed early.
 	/// 3. Wait for the block to be imported via subscription.
 	/// 4. If timeout is reached, we return an error.
-	async fn wait_for_block(&self, wait_for_hash: PHash) -> RelayChainResult<()> {
-		let mut head_stream = self.rpc_client.get_imported_heads_stream().await?;
+	async fn wait_for_block(&self, wait_for_hash: RelayHash) -> RelayChainResult<()> {
+		let mut head_stream = self.rpc_client.get_imported_heads_stream()?;

 		if self.rpc_client.chain_get_header(Some(wait_for_hash)).await?.is_some() {
 			return Ok(())
@@ -191,8 +192,8 @@ impl RelayChainInterface for RelayChainRpcInterface {

 	async fn new_best_notification_stream(
 		&self,
-	) -> RelayChainResult<Pin<Box<dyn Stream<Item = PHeader> + Send>>> {
-		let imported_headers_stream = self.rpc_client.get_best_heads_stream().await?;
+	) -> RelayChainResult<Pin<Box<dyn Stream<Item = RelayHeader> + Send>>> {
+		let imported_headers_stream = self.rpc_client.get_best_heads_stream()?;
 		Ok(imported_headers_stream.boxed())
 	}
 }
@@ -0,0 +1,520 @@
+// Copyright 2022 Parity Technologies (UK) Ltd.
+// This file is part of Cumulus.
+
+// Cumulus 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.
+
+// Cumulus 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 Cumulus.  If not, see <http://www.gnu.org/licenses/>.
+
+use cumulus_primitives_core::relay_chain::{
+	BlockNumber as RelayBlockNumber, Header as RelayHeader,
+};
+use cumulus_relay_chain_interface::{RelayChainError, RelayChainResult};
+use futures::{
+	channel::{
+		mpsc::{Receiver, Sender},
+		oneshot::Sender as OneshotSender,
+	},
+	future::BoxFuture,
+	stream::FuturesUnordered,
+	FutureExt, StreamExt,
+};
+use jsonrpsee::{
+	core::{
+		client::{Client as JsonRpcClient, ClientT, Subscription, SubscriptionClientT},
+		params::ArrayParams,
+		Error as JsonRpseeError, JsonValue,
+	},
+	rpc_params,
+	ws_client::WsClientBuilder,
+};
+use lru::LruCache;
+use polkadot_service::TaskManager;
+use std::{num::NonZeroUsize, sync::Arc};
+use tokio::sync::mpsc::{
+	channel as tokio_channel, Receiver as TokioReceiver, Sender as TokioSender,
+};
+use url::Url;
+
+const NOTIFICATION_CHANNEL_SIZE_LIMIT: usize = 20;
+const LOG_TARGET: &str = "reconnecting-websocket-client";
+
+/// Messages for communication between [`ReconnectingWsClient`] and [`ReconnectingWebsocketWorker`].
+#[derive(Debug)]
+enum RpcDispatcherMessage {
+	RegisterBestHeadListener(Sender<RelayHeader>),
+	RegisterImportListener(Sender<RelayHeader>),
+	RegisterFinalizationListener(Sender<RelayHeader>),
+	Request(String, ArrayParams, OneshotSender<Result<JsonValue, JsonRpseeError>>),
+}
+
+/// Frontend for performing websocket requests.
+/// Requests and stream requests are forwarded to [`ReconnectingWebsocketWorker`].
+#[derive(Debug, Clone)]
+pub struct ReconnectingWsClient {
+	/// Channel to communicate with the RPC worker
+	to_worker_channel: TokioSender<RpcDispatcherMessage>,
+}
+
+impl ReconnectingWsClient {
+	/// Create a new websocket client frontend.
+	pub async fn new(urls: Vec<Url>, task_manager: &mut TaskManager) -> RelayChainResult<Self> {
+		tracing::debug!(target: LOG_TARGET, "Instantiating reconnecting websocket client");
+
+		let (worker, sender) = ReconnectingWebsocketWorker::new(urls).await;
+
+		task_manager
+			.spawn_essential_handle()
+			.spawn("relay-chain-rpc-worker", None, worker.run());
+
+		Ok(Self { to_worker_channel: sender })
+	}
+}
+
+impl ReconnectingWsClient {
+	/// Perform a request via websocket connection.
+	pub async fn request<R>(&self, method: &str, params: ArrayParams) -> Result<R, RelayChainError>
+	where
+		R: serde::de::DeserializeOwned,
+	{
+		let (tx, rx) = futures::channel::oneshot::channel();
+
+		let message = RpcDispatcherMessage::Request(method.into(), params, tx);
+		self.to_worker_channel.send(message).await.map_err(|err| {
+			RelayChainError::WorkerCommunicationError(format!(
+				"Unable to send message to RPC worker: {}",
+				err
+			))
+		})?;
+
+		let value = rx.await.map_err(|err| {
+			RelayChainError::WorkerCommunicationError(format!(
+				"Unexpected channel close on RPC worker side: {}",
+				err
+			))
+		})??;
+
+		serde_json::from_value(value)
+			.map_err(|_| RelayChainError::GenericError("Unable to deserialize value".to_string()))
+	}
+
+	/// Get a stream of new best relay chain headers
+	pub fn get_best_heads_stream(&self) -> Result<Receiver<RelayHeader>, RelayChainError> {
+		let (tx, rx) =
+			futures::channel::mpsc::channel::<RelayHeader>(NOTIFICATION_CHANNEL_SIZE_LIMIT);
+		self.send_register_message_to_worker(RpcDispatcherMessage::RegisterBestHeadListener(tx))?;
+		Ok(rx)
+	}
+
+	/// Get a stream of finalized relay chain headers
+	pub fn get_finalized_heads_stream(&self) -> Result<Receiver<RelayHeader>, RelayChainError> {
+		let (tx, rx) =
+			futures::channel::mpsc::channel::<RelayHeader>(NOTIFICATION_CHANNEL_SIZE_LIMIT);
+		self.send_register_message_to_worker(RpcDispatcherMessage::RegisterFinalizationListener(
+			tx,
+		))?;
+		Ok(rx)
+	}
+
+	/// Get a stream of all imported relay chain headers
+	pub fn get_imported_heads_stream(&self) -> Result<Receiver<RelayHeader>, RelayChainError> {
+		let (tx, rx) =
+			futures::channel::mpsc::channel::<RelayHeader>(NOTIFICATION_CHANNEL_SIZE_LIMIT);
+		self.send_register_message_to_worker(RpcDispatcherMessage::RegisterImportListener(tx))?;
+		Ok(rx)
+	}
+
+	fn send_register_message_to_worker(
+		&self,
+		message: RpcDispatcherMessage,
+	) -> Result<(), RelayChainError> {
+		self.to_worker_channel
+			.try_send(message)
+			.map_err(|e| RelayChainError::WorkerCommunicationError(e.to_string()))
+	}
+}
+
+/// Worker that should be used in combination with [`RelayChainRpcClient`]. Must be polled to distribute header notifications to listeners.
+struct ReconnectingWebsocketWorker {
+	ws_urls: Vec<Url>,
+	/// Communication channel with the RPC client
+	client_receiver: TokioReceiver<RpcDispatcherMessage>,
+
+	/// Senders to distribute incoming header notifications to
+	imported_header_listeners: Vec<Sender<RelayHeader>>,
+	finalized_header_listeners: Vec<Sender<RelayHeader>>,
+	best_header_listeners: Vec<Sender<RelayHeader>>,
+}
+
+fn distribute_header(header: RelayHeader, senders: &mut Vec<Sender<RelayHeader>>) {
+	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,
+				}
+			});
+}
+
+/// Manages the active websocket client.
+/// Responsible for creating request futures, subscription streams
+/// and reconnections.
+#[derive(Debug)]
+struct ClientManager {
+	urls: Vec<Url>,
+	active_client: Arc<JsonRpcClient>,
+	active_index: usize,
+}
+
+struct RelayChainSubscriptions {
+	import_subscription: Subscription<RelayHeader>,
+	finalized_subscription: Subscription<RelayHeader>,
+	best_subscription: Subscription<RelayHeader>,
+}
+
+/// Try to find a new RPC server to connect to.
+async fn connect_next_available_rpc_server(
+	urls: &Vec<Url>,
+	starting_position: usize,
+) -> Result<(usize, Arc<JsonRpcClient>), ()> {
+	tracing::debug!(target: LOG_TARGET, starting_position, "Connecting to RPC server.");
+	for (counter, url) in urls.iter().cycle().skip(starting_position).take(urls.len()).enumerate() {
+		let index = (starting_position + counter) % urls.len();
+		tracing::info!(
+			target: LOG_TARGET,
+			index,
+			?url,
+			"Trying to connect to next external relaychain node.",
+		);
+		if let Ok(ws_client) = WsClientBuilder::default().build(url).await {
+			return Ok((index, Arc::new(ws_client)))
+		};
+	}
+	Err(())
+}
+
+impl ClientManager {
+	pub async fn new(urls: Vec<Url>) -> Result<Self, ()> {
+		if urls.is_empty() {
+			return Err(())
+		}
+		let active_client = connect_next_available_rpc_server(&urls, 0).await?;
+		Ok(Self { urls, active_client: active_client.1, active_index: active_client.0 })
+	}
+
+	pub async fn connect_to_new_rpc_server(&mut self) -> Result<(), ()> {
+		let new_active =
+			connect_next_available_rpc_server(&self.urls, self.active_index + 1).await?;
+		self.active_client = new_active.1;
+		self.active_index = new_active.0;
+		Ok(())
+	}
+
+	async fn get_subscriptions(&self) -> Result<RelayChainSubscriptions, JsonRpseeError> {
+		let import_subscription = self
+			.active_client
+			.subscribe::<RelayHeader, _>(
+				"chain_subscribeAllHeads",
+				rpc_params![],
+				"chain_unsubscribeAllHeads",
+			)
+			.await
+			.map_err(|e| {
+				tracing::error!(
+					target: LOG_TARGET,
+					?e,
+					"Unable to open `chain_subscribeAllHeads` subscription."
+				);
+				e
+			})?;
+		let best_subscription = self
+			.active_client
+			.subscribe::<RelayHeader, _>(
+				"chain_subscribeNewHeads",
+				rpc_params![],
+				"chain_unsubscribeNewHeads",
+			)
+			.await
+			.map_err(|e| {
+				tracing::error!(
+					target: LOG_TARGET,
+					?e,
+					"Unable to open `chain_subscribeNewHeads` subscription."
+				);
+				e
+			})?;
+		let finalized_subscription = self
+			.active_client
+			.subscribe::<RelayHeader, _>(
+				"chain_subscribeFinalizedHeads",
+				rpc_params![],
+				"chain_unsubscribeFinalizedHeads",
+			)
+			.await
+			.map_err(|e| {
+				tracing::error!(
+					target: LOG_TARGET,
+					?e,
+					"Unable to open `chain_subscribeFinalizedHeads` subscription."
+				);
+				e
+			})?;
+
+		Ok(RelayChainSubscriptions {
+			import_subscription,
+			best_subscription,
+			finalized_subscription,
+		})
+	}
+
+	/// Create a request future that performs an RPC request and sends the results to the caller.
+	/// In case of a dead websocket connection, it returns the original request parameters to
+	/// enable retries.
+	fn create_request(
+		&self,
+		method: String,
+		params: ArrayParams,
+		response_sender: OneshotSender<Result<JsonValue, JsonRpseeError>>,
+	) -> BoxFuture<'static, Result<(), RpcDispatcherMessage>> {
+		let future_client = self.active_client.clone();
+		async move {
+			let resp = future_client.request(&method, params.clone()).await;
+
+			// We should only return the original request in case
+			// the websocket connection is dead and requires a restart.
+			// Other errors should be forwarded to the request caller.
+			if let Err(JsonRpseeError::RestartNeeded(_)) = resp {
+				return Err(RpcDispatcherMessage::Request(method, params, response_sender))
+			}
+
+			if let Err(err) = response_sender.send(resp) {
+				tracing::debug!(
+					target: LOG_TARGET,
+					?err,
+					"Recipient no longer interested in request result"
+				);
+			}
+			Ok(())
+		}
+		.boxed()
+	}
+}
+
+enum ConnectionStatus {
+	Connected,
+	ReconnectRequired(Option<RpcDispatcherMessage>),
+}
+
+impl ReconnectingWebsocketWorker {
+	/// Create new worker. Returns the worker and a channel to register new listeners.
+	async fn new(
+		urls: Vec<Url>,
+	) -> (ReconnectingWebsocketWorker, TokioSender<RpcDispatcherMessage>) {
+		let (tx, rx) = tokio_channel(100);
+		let worker = ReconnectingWebsocketWorker {
+			ws_urls: urls,
+			client_receiver: rx,
+			imported_header_listeners: Vec::new(),
+			finalized_header_listeners: Vec::new(),
+			best_header_listeners: Vec::new(),
+		};
+		(worker, tx)
+	}
+
+	/// Reconnect via [`ClientManager`] and provide new notification streams.
+	async fn handle_reconnect(
+		&mut self,
+		client_manager: &mut ClientManager,
+		pending_requests: &mut FuturesUnordered<
+			BoxFuture<'static, Result<(), RpcDispatcherMessage>>,
+		>,
+		first_failed_request: Option<RpcDispatcherMessage>,
+	) -> Result<RelayChainSubscriptions, String> {
+		let mut requests_to_retry = Vec::new();
+		if let Some(req @ RpcDispatcherMessage::Request(_, _, _)) = first_failed_request {
+			requests_to_retry.push(req);
+		}
+
+		// At this point, all pending requests will return an error since the
+		// websocket connection is dead. So draining the pending requests should be fast.
+		while !pending_requests.is_empty() {
+			if let Some(Err(req)) = pending_requests.next().await {
+				requests_to_retry.push(req);
+			}
+		}
+
+		if client_manager.connect_to_new_rpc_server().await.is_err() {
+			return Err(format!("Unable to find valid external RPC server, shutting down."))
+		};
+
+		for item in requests_to_retry.into_iter() {
+			if let RpcDispatcherMessage::Request(method, params, response_sender) = item {
+				pending_requests.push(client_manager.create_request(
+					method,
+					params,
+					response_sender,
+				));
+			};
+		}
+
+		client_manager.get_subscriptions().await.map_err(|e| {
+			format!("Not able to create streams from newly connected RPC server, shutting down. err: {:?}", e)
+		})
+	}
+
+	/// Run this worker to drive notification streams.
+	/// The worker does the following:
+	/// - Listen for [`RpcDispatcherMessage`], perform requests and register new listeners for the notification streams
+	/// - 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.
+	/// - Find a new valid RPC server to connect to in case the websocket connection is terminated.
+	///   If the worker is not able to connec to an RPC server from the list, the worker shuts down.
+	async fn run(mut self) {
+		let mut pending_requests = FuturesUnordered::new();
+
+		let urls = std::mem::take(&mut self.ws_urls);
+		let Ok(mut client_manager) = ClientManager::new(urls).await else {
+			tracing::error!(target: LOG_TARGET, "No valid RPC url found. Stopping RPC worker.");
+			return;
+		};
+		let Ok(mut subscriptions) = client_manager.get_subscriptions().await else {
+			tracing::error!(target: LOG_TARGET, "Unable to fetch subscriptions on initial connection.");
+			return;
+		};
+
+		let mut imported_blocks_cache =
+			LruCache::new(NonZeroUsize::new(40).expect("40 is nonzero; qed."));
+		let mut should_reconnect = ConnectionStatus::Connected;
+		let mut last_seen_finalized_num: RelayBlockNumber = 0;
+		loop {
+			// This branch is taken if the websocket connection to the current RPC server is closed.
+			if let ConnectionStatus::ReconnectRequired(maybe_failed_request) = should_reconnect {
+				match self
+					.handle_reconnect(
+						&mut client_manager,
+						&mut pending_requests,
+						maybe_failed_request,
+					)
+					.await
+				{
+					Ok(new_subscriptions) => {
+						subscriptions = new_subscriptions;
+					},
+					Err(message) => {
+						tracing::error!(
+							target: LOG_TARGET,
+							message,
+							"Unable to reconnect, stopping worker."
+						);
+						return
+					},
+				}
+				should_reconnect = ConnectionStatus::Connected;
+			}
+
+			tokio::select! {
+				evt = self.client_receiver.recv() => match evt {
+					Some(RpcDispatcherMessage::RegisterBestHeadListener(tx)) => {
+						self.best_header_listeners.push(tx);
+					},
+					Some(RpcDispatcherMessage::RegisterImportListener(tx)) => {
+						self.imported_header_listeners.push(tx)
+					},
+					Some(RpcDispatcherMessage::RegisterFinalizationListener(tx)) => {
+						self.finalized_header_listeners.push(tx)
+					},
+					Some(RpcDispatcherMessage::Request(method, params, response_sender)) => {
+						pending_requests.push(client_manager.create_request(method, params, response_sender));
+					},
+					None => {
+						tracing::error!(target: LOG_TARGET, "RPC client receiver closed. Stopping RPC Worker.");
+						return;
+					}
+				},
+				should_retry = pending_requests.next(), if !pending_requests.is_empty() => {
+					if let Some(Err(req)) = should_retry {
+						should_reconnect = ConnectionStatus::ReconnectRequired(Some(req));
+					}
+				},
+				import_event = subscriptions.import_subscription.next() => {
+					match import_event {
+						Some(Ok(header)) => {
+							let hash = header.hash();
+							if imported_blocks_cache.contains(&hash) {
+								tracing::debug!(
+									target: LOG_TARGET,
+									number = header.number,
+									?hash,
+									"Duplicate imported block header. This might happen after switching to a new RPC node. Skipping distribution."
+								);
+								continue;
+							}
+							imported_blocks_cache.put(hash, ());
+							distribute_header(header, &mut self.imported_header_listeners);
+						},
+						None => {
+							tracing::error!(target: LOG_TARGET, "Subscription closed.");
+							should_reconnect = ConnectionStatus::ReconnectRequired(None);
+						},
+						Some(Err(error)) => {
+							tracing::error!(target: LOG_TARGET, ?error, "Error in RPC subscription.");
+							should_reconnect = ConnectionStatus::ReconnectRequired(None);
+						},
+					}
+				},
+				best_header_event = subscriptions.best_subscription.next() => {
+					match best_header_event {
+						Some(Ok(header)) => distribute_header(header, &mut self.best_header_listeners),
+						None => {
+							tracing::error!(target: LOG_TARGET, "Subscription closed.");
+							should_reconnect = ConnectionStatus::ReconnectRequired(None);
+						},
+						Some(Err(error)) => {
+							tracing::error!(target: LOG_TARGET, ?error, "Error in RPC subscription.");
+							should_reconnect = ConnectionStatus::ReconnectRequired(None);
+						},
+					}
+				}
+				finalized_event = subscriptions.finalized_subscription.next() => {
+					match finalized_event {
+						Some(Ok(header)) if header.number > last_seen_finalized_num => {
+							last_seen_finalized_num = header.number;
+							distribute_header(header, &mut self.finalized_header_listeners);
+						},
+						Some(Ok(header)) => {
+							tracing::debug!(
+								target: LOG_TARGET,
+								number = header.number,
+								last_seen_finalized_num,
+								"Duplicate finalized block header. This might happen after switching to a new RPC node. Skipping distribution."
+							);
+						},
+						None => {
+							tracing::error!(target: LOG_TARGET, "Subscription closed.");
+							should_reconnect = ConnectionStatus::ReconnectRequired(None);
+						},
+						Some(Err(error)) => {
+							tracing::error!(target: LOG_TARGET, ?error, "Error in RPC subscription.");
+							should_reconnect = ConnectionStatus::ReconnectRequired(None);
+						},
+					}
+				}
+			}
+		}
+	}
+}
@@ -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()
 	}
 }