Implement ping-pong for WebSocket clients

client/transport/src/ws/mod.rs

@@ -31,16 +31,17 @@ use std::net::{SocketAddr, ToSocketAddrs};
 use std::time::Duration;
 
 use futures_util::io::{BufReader, BufWriter};
-use jsonrpsee_core::client::{CertificateStore, TransportReceiverT, TransportSenderT};
+use jsonrpsee_core::client::{CertificateStore, ReceivedMessage, TransportReceiverT, TransportSenderT};
 use jsonrpsee_core::TEN_MB_SIZE_BYTES;
 use jsonrpsee_core::{async_trait, Cow};
-use soketto::connection;
 use soketto::handshake::client::{Client as WsHandshakeClient, ServerResponse};
+use soketto::{connection, Incoming};
 use stream::EitherStream;
 use thiserror::Error;
 use tokio::net::TcpStream;
 
 pub use http::{uri::InvalidUri, Uri};
+use soketto::data::ByteSlice125;
 pub use soketto::handshake::client::Header;
 
 /// Sending end of WebSocket transport.
@@ -195,6 +196,16 @@ impl TransportSenderT for Sender {
 		Ok(())
 	}
 
+	/// Sends out a ping request. Returns a `Future` that finishes when the request has been
+	/// successfully sent.
+	async fn send_ping(&mut self, data: &[u8]) -> Result<(), Self::Error> {
+		tracing::debug!("send ping: {:?}", data);
+		let byte_slice = ByteSlice125::try_from(data).expect("Found invalid ping slice");
+		self.inner.send_ping(byte_slice).await?;
+		self.inner.flush().await?;
+		Ok(())
+	}
+
 	/// Send a close message and close the connection.
 	async fn close(&mut self) -> Result<(), WsError> {
 		self.inner.close().await.map_err(Into::into)
@@ -206,11 +217,19 @@ impl TransportReceiverT for Receiver {
 	type Error = WsError;
 
 	/// Returns a `Future` resolving when the server sent us something back.
-	async fn receive(&mut self) -> Result<String, Self::Error> {
+	async fn receive(&mut self) -> Result<ReceivedMessage, Self::Error> {
 		let mut message = Vec::new();
-		self.inner.receive_data(&mut message).await?;
-		let s = String::from_utf8(message).expect("Found invalid UTF-8");
-		Ok(s)
+
+		loop {
+			let recv = self.inner.receive(&mut message).await?;
+
+			if let Incoming::Data(_) = recv {
+				let s = String::from_utf8(message).expect("Found invalid UTF-8");
+				return Ok(ReceivedMessage::Data(s));
+			} else if let Incoming::Pong(pong_data) = recv {
+				return Ok(ReceivedMessage::Pong(Vec::from(pong_data)));
+			}
+		}
 	}
 }
 

client/ws-client/src/lib.rs

@@ -73,6 +73,7 @@ pub struct WsClientBuilder<'a> {
 	max_request_body_size: u32,
 	request_timeout: Duration,
 	connection_timeout: Duration,
+	ping_interval: Option<Duration>,
 	headers: Vec<Header<'a>>,
 	max_concurrent_requests: usize,
 	max_notifs_per_subscription: usize,
@@ -87,6 +88,7 @@ impl<'a> Default for WsClientBuilder<'a> {
 			max_request_body_size: TEN_MB_SIZE_BYTES,
 			request_timeout: Duration::from_secs(60),
 			connection_timeout: Duration::from_secs(10),
+			ping_interval: None,
 			headers: Vec::new(),
 			max_concurrent_requests: 256,
 			max_notifs_per_subscription: 1024,
@@ -121,6 +123,12 @@ impl<'a> WsClientBuilder<'a> {
 		self
 	}
 
+	/// See documentation [`ClientBuilder::ping_interval`] (disabled by default).
+	pub fn ping_interval(mut self, interval: Duration) -> Self {
+		self.ping_interval = Some(interval);
+		self
+	}
+
 	/// See documentation [`WsTransportClientBuilder::add_header`] (default is none).
 	pub fn add_header(mut self, name: &'a str, value: &'a str) -> Self {
 		self.headers.push(Header { name, value: value.as_bytes() });
@@ -169,11 +177,16 @@ impl<'a> WsClientBuilder<'a> {
 		let uri: Uri = url.as_ref().parse().map_err(|e: InvalidUri| Error::Transport(e.into()))?;
 		let (sender, receiver) = transport_builder.build(uri).await.map_err(|e| Error::Transport(e.into()))?;
 
-		Ok(ClientBuilder::default()
+		let mut client = ClientBuilder::default()
 			.max_notifs_per_subscription(self.max_notifs_per_subscription)
 			.request_timeout(self.request_timeout)
 			.max_concurrent_requests(self.max_concurrent_requests)
-			.id_format(self.id_kind)
-			.build_with_tokio(sender, receiver))
+			.id_format(self.id_kind);
+
+		if let Some(interval) = self.ping_interval {
+			client = client.ping_interval(interval);
+		}
+
+		Ok(client.build_with_tokio(sender, receiver))
 	}
 }