fix(server): remove MethodSinkPermit to fix backpressure issue on c…

…oncurrent subscriptions (paritytech#1126)

* fix(server): don't use `Permit` for responses

After a method call has been received, this PR doesn't keep the `Permit`
anymore and requires each response to await for a slot in the mpsc buffer.

This is especially important for subscriptions which previously
two slots were unintentionally required because the server kept one around until
the subscription was "accepted" and the subscription itself waited for a slot.

If many concurrent subscriptions were received "concurrently" and "filled" the buffer
no subscription could make progress.

* cleanup code

* Update server/src/transport/ws.rs

Co-authored-by: James Wilson <james@jsdw.me>

* Update core/src/server/helpers.rs

Co-authored-by: James Wilson <james@jsdw.me>

* cargo fmt

---------

Co-authored-by: James Wilson <james@jsdw.me>

core/src/server/helpers.rs

@@ -34,7 +34,7 @@ use jsonrpsee_types::error::{
 use jsonrpsee_types::{Id, InvalidRequest, Response, ResponsePayload};
 use serde::Serialize;
 use serde_json::value::to_raw_value;
-use tokio::sync::mpsc::{self, OwnedPermit};
+use tokio::sync::mpsc;
 
 use super::{DisconnectError, SendTimeoutError, SubscriptionMessage, TrySendError};
 
@@ -139,45 +139,31 @@ impl MethodSink {
 		self.tx.send(msg).await.map_err(Into::into)
 	}
 
+	/// Send a JSON-RPC error to the client
+	pub async fn send_error<'a>(&self, id: Id<'a>, err: ErrorObject<'a>) -> Result<(), DisconnectError> {
+		let json =
+			serde_json::to_string(&Response::new(ResponsePayload::<()>::Error(err), id)).expect("valid JSON; qed");
+
+		self.send(json).await
+	}
+
 	/// Similar to to `MethodSink::send` but only waits for a limited time.
 	pub async fn send_timeout(&self, msg: String, timeout: Duration) -> Result<(), SendTimeoutError> {
 		tx_log_from_str(&msg, self.max_log_length);
 		self.tx.send_timeout(msg, timeout).await.map_err(Into::into)
 	}
 
-	/// Waits for channel capacity. Once capacity to send one message is available, it is reserved for the caller.
-	pub async fn reserve(&self) -> Result<MethodSinkPermit, DisconnectError> {
-		match self.tx.clone().reserve_owned().await {
-			Ok(permit) => Ok(MethodSinkPermit { tx: permit, max_log_length: self.max_log_length }),
+	/// Waits for there to be space on the return channel.
+	pub async fn has_capacity(&self) -> Result<(), DisconnectError> {
+		match self.tx.reserve().await {
+			// The permit is thrown away here because it's just
+			// a way to ensure that the return buffer has space.
+			Ok(_) => Ok(()),
 			Err(_) => Err(DisconnectError(SubscriptionMessage::empty())),
 		}
 	}
 }
 
-/// A method sink with reserved spot in the bounded queue.
-#[derive(Debug)]
-pub struct MethodSinkPermit {
-	tx: OwnedPermit<String>,
-	max_log_length: u32,
-}
-
-impl MethodSinkPermit {
-	/// Send a JSON-RPC error to the client
-	pub fn send_error(self, id: Id, err: ErrorObject) {
-		let json = serde_json::to_string(&Response::new(ResponsePayload::<()>::Error(err.into_owned()), id))
-			.expect("valid JSON; qed");
-
-		self.send_raw(json)
-	}
-
-	/// Send a raw JSON-RPC message to the client, `MethodSink` does not check the validity
-	/// of the JSON being sent.
-	pub fn send_raw(self, json: String) {
-		self.tx.send(json.clone());
-		tx_log_from_str(&json, self.max_log_length);
-	}
-}
-
 /// Figure out if this is a sufficiently complete request that we can extract an [`Id`] out of, or just plain
 /// unparseable garbage.
 pub fn prepare_error(data: &[u8]) -> (Id<'_>, ErrorCode) {

core/src/server/mod.rs

@@ -38,7 +38,7 @@ mod rpc_module;
 mod subscription;
 
 pub use error::*;
-pub use helpers::{BatchResponseBuilder, BoundedWriter, MethodResponse, MethodSink, MethodSinkPermit};
+pub use helpers::{BatchResponseBuilder, BoundedWriter, MethodResponse, MethodSink};
 pub use host_filtering::*;
 pub use rpc_module::*;
 pub use subscription::*;

server/src/transport/ws.rs

@@ -12,9 +12,7 @@ use hyper::upgrade::Upgraded;
 use jsonrpsee_core::server::helpers::{
 	batch_response_error, prepare_error, BatchResponseBuilder, MethodResponse, MethodSink,
 };
-use jsonrpsee_core::server::{
-	BoundedSubscriptions, CallOrSubscription, MethodCallback, MethodSinkPermit, Methods, SubscriptionState,
-};
+use jsonrpsee_core::server::{BoundedSubscriptions, CallOrSubscription, MethodCallback, Methods, SubscriptionState};
 use jsonrpsee_core::tracing::{rx_log_from_json, tx_log_from_str};
 use jsonrpsee_core::traits::IdProvider;
 use jsonrpsee_core::{Error, JsonRawValue};
@@ -263,14 +261,19 @@ pub(crate) async fn background_task<L: Logger>(sender: Sender, mut receiver: Rec
 	let result = loop {
 		data.clear();
 
-		let sink_permit = match wait_for_permit(&sink, stopped).await {
-			Some((permit, stop)) => {
-				stopped = stop;
-				permit
-			}
-			None => break Ok(Shutdown::ConnectionClosed),
+		// This is a guard to ensure that the underlying socket is only read if there is space in
+		// the buffer for messages to be sent back to them.
+		//
+		// Thus, this check enforces that if the client can't keep up with receiving messages,
+		// then no new messages will be read from them.
+		//
+		// TCP retransmission mechanism will take of the rest and adjust the window size accordingly.
+		let Some(stop) = wait_until_connection_buffer_has_capacity(&sink, stopped).await else {
+			break Ok(Shutdown::ConnectionClosed)
 		};
 
+		stopped = stop;
+
 		match try_recv(&mut receiver, &mut data, stopped).await {
 			Receive::Shutdown => break Ok(Shutdown::Stopped),
 			Receive::Ok(stop) => {
@@ -290,7 +293,9 @@ pub(crate) async fn background_task<L: Logger>(sender: Sender, mut receiver: Rec
 							current,
 							maximum
 						);
-						sink_permit.send_error(Id::Null, reject_too_big_request(max_request_body_size));
+						if sink.send_error(Id::Null, reject_too_big_request(max_request_body_size)).await.is_err() {
+							break Ok(Shutdown::ConnectionClosed);
+						}
 
 						continue;
 					}
@@ -310,7 +315,6 @@ pub(crate) async fn background_task<L: Logger>(sender: Sender, mut receiver: Rec
 			max_log_length,
 			max_response_body_size,
 			sink: sink.clone(),
-			sink_permit,
 			id_provider: id_provider.clone(),
 			logger: logger.clone(),
 			data: std::mem::take(&mut data),
@@ -403,21 +407,18 @@ enum Receive<S> {
 	Ok(S),
 }
 
-// Wait until there is a slot in the bounded channel.
-//
-// This will force the client to read socket on the other side
-// otherwise the socket will not be read again.
+// Wait until there is capacity in connection buffer to send one message.
 //
 // Fails if the server was stopped.
-async fn wait_for_permit<S>(sink: &MethodSink, stopped: S) -> Option<(MethodSinkPermit, S)>
+async fn wait_until_connection_buffer_has_capacity<S>(sink: &MethodSink, stopped: S) -> Option<S>
 where
 	S: Future<Output = ()> + Unpin,
 {
-	let reserve = sink.reserve();
+	let reserve = sink.has_capacity();
 	tokio::pin!(reserve);
 
 	match futures_util::future::select(reserve, stopped).await {
-		Either::Left((Ok(sink), s)) => Some((sink, s)),
+		Either::Left((Ok(_), s)) => Some(s),
 		_ => None,
 	}
 }
@@ -461,7 +462,6 @@ struct ExecuteCallParams<L: Logger> {
 	max_response_body_size: u32,
 	max_log_length: u32,
 	sink: MethodSink,
-	sink_permit: MethodSinkPermit,
 	logger: L,
 }
 
@@ -471,7 +471,6 @@ async fn execute_unchecked_call<L: Logger>(params: ExecuteCallParams<L>) {
 		conn_id,
 		data,
 		sink,
-		sink_permit,
 		max_response_body_size,
 		max_log_length,
 		methods,
@@ -505,7 +504,7 @@ async fn execute_unchecked_call<L: Logger>(params: ExecuteCallParams<L>) {
 
 					CallOrSubscription::Call(r) => {
 						logger.on_response(&r.result, request_start, TransportProtocol::WebSocket);
-						sink_permit.send_raw(r.result);
+						_ = sink.send(r.result).await;
 					}
 				}
 			}
@@ -518,7 +517,7 @@ async fn execute_unchecked_call<L: Logger>(params: ExecuteCallParams<L>) {
 						ErrorObject::borrowed(BATCHES_NOT_SUPPORTED_CODE, &BATCHES_NOT_SUPPORTED_MSG, None),
 					);
 					logger.on_response(&response.result, request_start, TransportProtocol::WebSocket);
-					sink_permit.send_raw(response.result);
+					_ = sink.send(response.result).await;
 					return;
 				}
 				BatchRequestConfig::Limit(limit) => limit as usize,
@@ -542,11 +541,11 @@ async fn execute_unchecked_call<L: Logger>(params: ExecuteCallParams<L>) {
 			if let Some(response) = response {
 				tx_log_from_str(&response, max_log_length);
 				logger.on_response(&response, request_start, TransportProtocol::WebSocket);
-				sink_permit.send_raw(response);
+				_ = sink.send(response).await;
 			}
 		}
 		_ => {
-			sink_permit.send_error(Id::Null, ErrorCode::ParseError.into());
+			_ = sink.send_error(Id::Null, ErrorCode::ParseError.into()).await;
 		}
 	};
 }