Implement ping-pong for WebSocket clients

client/transport/src/web.rs

@@ -5,7 +5,7 @@ use futures_util::sink::SinkExt;
 use futures_util::stream::{SplitSink, SplitStream, StreamExt};
 use gloo_net::websocket::{futures::WebSocket, Message, WebSocketError};
 use jsonrpsee_core::async_trait;
-use jsonrpsee_core::client::{TransportReceiverT, TransportSenderT};
+use jsonrpsee_core::client::{ReceivedMessage, TransportReceiverT, TransportSenderT};
 
 /// Web-sys transport error that can occur.
 #[derive(Debug, thiserror::Error)]
@@ -22,6 +22,9 @@ pub enum Error {
 	/// WebSocket error
 	#[error("WebSocket Error: {0:?}")]
 	WebSocket(WebSocketError),
+	/// Operation not supported
+	#[error("Operation not supported")]
+	NotSupported,
 }
 
 /// Sender.
@@ -52,6 +55,11 @@ impl TransportSenderT for Sender {
 		Ok(())
 	}
 
+	async fn send_ping(&mut self) -> Result<(), Self::Error> {
+		tracing::trace!("send ping - not implemented for wasm");
+		Err(Error::NotSupported)
+	}
+
 	async fn close(&mut self) -> Result<(), Error> {
 		Ok(())
 	}
@@ -61,17 +69,15 @@ impl TransportSenderT for Sender {
 impl TransportReceiverT for Receiver {
 	type Error = Error;
 
-	async fn receive(&mut self) -> Result<String, Self::Error> {
+	async fn receive(&mut self) -> Result<ReceivedMessage, Self::Error> {
 		match self.0.next().await {
 			Some(Ok(msg)) => {
 				tracing::trace!("rx: {:?}", msg);
 
-				let txt = match msg {
-					Message::Bytes(bytes) => String::from_utf8(bytes).expect("WebSocket message is valid utf8; qed"),
-					Message::Text(txt) => txt,
-				};
-
-				Ok(txt)
+				match msg {
+					Message::Bytes(bytes) => Ok(ReceivedMessage::Bytes(bytes)),
+					Message::Text(txt) => Ok(ReceivedMessage::Text(txt)),
+				}
 			}
 			Some(Err(err)) => Err(Error::WebSocket(err)),
 			None => Err(Error::SenderDisconnected),

client/transport/src/ws/mod.rs

@@ -31,11 +31,13 @@ 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::connection::Error::Utf8;
+use soketto::data::ByteSlice125;
 use soketto::handshake::client::{Client as WsHandshakeClient, ServerResponse};
+use soketto::{connection, Data, Incoming};
 use stream::EitherStream;
 use thiserror::Error;
 use tokio::net::TcpStream;
@@ -195,6 +197,20 @@ 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) -> Result<(), Self::Error> {
+		tracing::debug!("send ping");
+		// Submit empty slice as "optional" parameter.
+		let slice: &[u8] = &[];
+		// Byte slice fails if the provided slice is larger than 125 bytes.
+		let byte_slice = ByteSlice125::try_from(slice).expect("Empty slice should fit into ByteSlice125");
+
+		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 +222,21 @@ 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> {
-		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)
+	async fn receive(&mut self) -> Result<ReceivedMessage, Self::Error> {
+		loop {
+			let mut message = Vec::new();
+			let recv = self.inner.receive(&mut message).await?;
+
+			match recv {
+				Incoming::Data(Data::Text(_)) => {
+					let s = String::from_utf8(message).map_err(|err| WsError::Connection(Utf8(err.utf8_error())))?;
+					break Ok(ReceivedMessage::Text(s));
+				}
+				Incoming::Data(Data::Binary(_)) => break Ok(ReceivedMessage::Bytes(message)),
+				Incoming::Pong(_) => break Ok(ReceivedMessage::Pong),
+				_ => continue,
+			}
+		}
 	}
 }
 

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))
 	}
 }