client: Implement `notify_on_disconnect`

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>

core/src/client/async_client/mod.rs

@@ -11,6 +11,7 @@ use crate::client::{
 use crate::tracing::{rx_log_from_json, tx_log_from_str, RpcTracing};
 
 use core::time::Duration;
+use std::sync::Arc;
 use helpers::{
 	build_unsubscribe_message, call_with_timeout, process_batch_response, process_error_response, process_notification,
 	process_single_response, process_subscription_response, stop_subscription,
@@ -31,6 +32,7 @@ use jsonrpsee_types::{
 	SubscriptionResponse,
 };
 use serde::de::DeserializeOwned;
+use tokio::sync::Notify;
 use tracing_futures::Instrument;
 
 use super::{FrontToBack, IdKind, RequestIdManager};
@@ -162,9 +164,20 @@ impl ClientBuilder {
 		let (err_tx, err_rx) = oneshot::channel();
 		let max_notifs_per_subscription = self.max_notifs_per_subscription;
 		let ping_interval = self.ping_interval;
+		let notify_on_close = Arc::new(Notify::new());
+		let notify_on_close_back = notify_on_close.clone();
 
 		tokio::spawn(async move {
-			background_task(sender, receiver, from_front, err_tx, max_notifs_per_subscription, ping_interval).await;
+			background_task(
+				sender,
+				receiver,
+				from_front,
+				err_tx,
+				max_notifs_per_subscription,
+				ping_interval,
+				notify_on_close_back,
+			)
+			.await;
 		});
 		Client {
 			to_back,
@@ -172,6 +185,7 @@ impl ClientBuilder {
 			error: Mutex::new(ErrorFromBack::Unread(err_rx)),
 			id_manager: RequestIdManager::new(self.max_concurrent_requests, self.id_kind),
 			max_log_length: self.max_log_length,
+			notify: notify_on_close,
 		}
 	}
 
@@ -216,6 +230,8 @@ pub struct Client {
 	///
 	/// Entries bigger than this limit will be truncated.
 	max_log_length: u32,
+	/// Notify when the client is disconnected or encountered an error.
+	notify: Arc<Notify>,
 }
 
 impl Client {
@@ -232,6 +248,16 @@ impl Client {
 		*err_lock = next_state;
 		err
 	}
+
+	/// Completes when the client is disconnected or the client's background task encountered an error.
+	///
+	/// # Cancel safety
+	///
+	/// This method is cancel safe. Once the client is disconnected, it stays disconnected forever and all
+	/// future calls to this method will return immediately.
+	pub async fn notify_on_disconnect(&self) {
+		self.notify.notified().await;
+	}
 }
 
 impl Drop for Client {
@@ -621,6 +647,7 @@ async fn background_task<S, R>(
 	front_error: oneshot::Sender<Error>,
 	max_notifs_per_subscription: usize,
 	ping_interval: Option<Duration>,
+	notify_on_close: Arc<Notify>,
 ) where
 	S: TransportSenderT,
 	R: TransportReceiverT,
@@ -692,6 +719,7 @@ async fn background_task<S, R>(
 			}
 		};
 	}
+	notify_on_close.notify_waiters();
 	// Send close message to the server.
 	let _ = sender.close().await;
 }