server: bounded channels and backpressure

benches/helpers.rs

@@ -146,9 +146,13 @@ pub async fn ws_server(handle: tokio::runtime::Handle) -> (String, jsonrpsee::se
 	let mut module = gen_rpc_module();
 
 	module
-		.register_subscription(SUB_METHOD_NAME, SUB_METHOD_NAME, UNSUB_METHOD_NAME, |_params, mut sink, _ctx| {
+		.register_subscription(SUB_METHOD_NAME, SUB_METHOD_NAME, UNSUB_METHOD_NAME, |_params, pending, _ctx| {
 			let x = "Hello";
-			tokio::spawn(async move { sink.send(&x) });
+			tokio::spawn(async move {
+				let sink = pending.accept().await.unwrap();
+				let msg = sink.build_message(&x).unwrap();
+				sink.send(msg).await.unwrap();
+			});
 			Ok(())
 		})
 		.unwrap();

core/Cargo.toml

@@ -45,7 +45,7 @@ server = [
 	"tokio/rt",
 	"tokio/sync",
 ]
-client = ["futures-util/sink", "futures-channel/sink", "futures-channel/std"]
+client = ["futures-util/sink", "futures-channel/sink"]
 async-client = [
 	"async-lock",
 	"client",

core/src/error.rs

@@ -111,21 +111,6 @@ pub enum Error {
 	/// Access control verification of HTTP headers failed.
 	#[error("HTTP header: `{0}` value: `{1}` verification failed")]
 	HttpHeaderRejected(&'static str, String),
-	/// Failed to execute a method because a resource was already at capacity
-	#[error("Resource at capacity: {0}")]
-	ResourceAtCapacity(&'static str),
-	/// Failed to register a resource due to a name conflict
-	#[error("Resource name already taken: {0}")]
-	ResourceNameAlreadyTaken(&'static str),
-	/// Failed to initialize resources for a method at startup
-	#[error("Resource name `{0}` not found for method `{1}`")]
-	ResourceNameNotFoundForMethod(&'static str, &'static str),
-	/// Trying to claim resources for a method execution, but the method resources have not been initialized
-	#[error("Method `{0}` has uninitialized resources")]
-	UninitializedMethod(Box<str>),
-	/// Failed to register a resource due to a maximum number of resources already registered
-	#[error("Maximum number of resources reached")]
-	MaxResourcesReached,
 	/// Custom error.
 	#[error("Custom error: {0}")]
 	Custom(String),

core/src/server/helpers.rs

@@ -25,15 +25,15 @@
 // DEALINGS IN THE SOFTWARE.
 
 use std::io;
-use std::sync::Arc;
 
 use crate::tracing::tx_log_from_str;
 use crate::Error;
-use futures_channel::mpsc;
 use jsonrpsee_types::error::{ErrorCode, ErrorObject, ErrorResponse, OVERSIZED_RESPONSE_CODE, OVERSIZED_RESPONSE_MSG};
 use jsonrpsee_types::{Id, InvalidRequest, Response};
 use serde::Serialize;
-use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore};
+use tokio::sync::mpsc::{self, Permit};
+
+use super::rpc_module::{DisconnectError, TrySendError};
 
 /// Bounded writer that allows writing at most `max_len` bytes.
 ///
@@ -84,7 +84,7 @@ impl<'a> io::Write for &'a mut BoundedWriter {
 #[derive(Clone, Debug)]
 pub struct MethodSink {
 	/// Channel sender.
-	tx: mpsc::UnboundedSender<String>,
+	pub tx: mpsc::Sender<String>,
 	/// Max response size in bytes for a executed call.
 	max_response_size: u32,
 	/// Max log length.
@@ -93,12 +93,12 @@ pub struct MethodSink {
 
 impl MethodSink {
 	/// Create a new `MethodSink` with unlimited response size.
-	pub fn new(tx: mpsc::UnboundedSender<String>) -> Self {
+	pub fn new(tx: mpsc::Sender<String>) -> Self {
 		MethodSink { tx, max_response_size: u32::MAX, max_log_length: u32::MAX }
 	}
 
 	/// Create a new `MethodSink` with a limited response size.
-	pub fn new_with_limit(tx: mpsc::UnboundedSender<String>, max_response_size: u32, max_log_length: u32) -> Self {
+	pub fn new_with_limit(tx: mpsc::Sender<String>, max_response_size: u32, max_log_length: u32) -> Self {
 		MethodSink { tx, max_response_size, max_log_length }
 	}
 
@@ -107,47 +107,68 @@ impl MethodSink {
 		self.tx.is_closed()
 	}
 
-	/// Send a JSON-RPC error to the client
-	pub fn send_error(&self, id: Id, error: ErrorObject) -> bool {
-		let json = match serde_json::to_string(&ErrorResponse::borrowed(error, id)) {
-			Ok(json) => json,
-			Err(err) => {
-				tracing::error!("Error serializing response: {:?}", err);
+	/// Same as [`tokio::sync::mpsc::Sender::closed`].
+	///
+	/// # Cancel safety
+	/// This method is cancel safe. Once the channel is closed,
+	/// it stays closed forever and all future calls to closed will return immediately.
+	pub async fn closed(&self) {
+		self.tx.closed().await
+	}
 
-				return false;
-			}
-		};
+	/// Get the max response size.
+	pub const fn max_response_size(&self) -> u32 {
+		self.max_response_size
+	}
 
-		tx_log_from_str(&json, self.max_log_length);
+	/// Non-blocking send which fails if the channel is closed or full
+	///
+	/// Returns the message if the send fails such that either can be thrown away or re-sent later.
+	pub fn try_send(&mut self, msg: String) -> Result<(), TrySendError> {
+		tx_log_from_str(&msg, self.max_log_length);
+		self.tx.try_send(msg)
+	}
 
-		if let Err(err) = self.send_raw(json) {
-			tracing::warn!("Error sending response {:?}", err);
-			false
-		} else {
-			true
+	/// Async send which will wait until there is space in channel buffer or that the subscription is disconnected.
+	pub async fn send(&self, msg: String) -> Result<(), DisconnectError<String>> {
+		tx_log_from_str(&msg, self.max_log_length);
+		self.tx.send(msg).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.reserve().await {
+			Ok(permit) => Ok(MethodSinkPermit { tx: permit, max_log_length: self.max_log_length }),
+			Err(e) => Err(e.into()),
 		}
 	}
+}
+
+/// A method sink with reserved spot in the bounded queue.
+#[derive(Debug)]
+pub struct MethodSinkPermit<'a> {
+	tx: Permit<'a, String>,
+	max_log_length: u32,
+}
+
+impl<'a> MethodSinkPermit<'a> {
+	/// Send a JSON-RPC error to the client
+	pub fn send_error(self, id: Id, error: ErrorObject) {
+		let json = serde_json::to_string(&ErrorResponse::borrowed(error, id)).expect("valid JSON; qed");
+
+		self.send_raw(json)
+	}
 
 	/// Helper for sending the general purpose `Error` as a JSON-RPC errors to the client.
-	pub fn send_call_error(&self, id: Id, err: Error) -> bool {
+	pub fn send_call_error(self, id: Id, err: Error) {
 		self.send_error(id, err.into())
 	}
 
 	/// Send a raw JSON-RPC message to the client, `MethodSink` does not check verify the validity
 	/// of the JSON being sent.
-	pub fn send_raw(&self, json: String) -> Result<(), mpsc::TrySendError<String>> {
+	pub fn send_raw(self, json: String) {
+		self.tx.send(json.clone());
 		tx_log_from_str(&json, self.max_log_length);
-		self.tx.unbounded_send(json)
-	}
-
-	/// Close the channel for any further messages.
-	pub fn close(&self) {
-		self.tx.close_channel();
-	}
-
-	/// Get the maximum number of permitted subscriptions.
-	pub const fn max_response_size(&self) -> u32 {
-		self.max_response_size
 	}
 }
 
@@ -160,59 +181,6 @@ pub fn prepare_error(data: &[u8]) -> (Id<'_>, ErrorCode) {
 	}
 }
 
-/// A permitted subscription.
-#[derive(Debug)]
-pub struct SubscriptionPermit {
-	_permit: OwnedSemaphorePermit,
-	resource: Arc<Notify>,
-}
-
-impl SubscriptionPermit {
-	/// Get the handle to [`tokio::sync::Notify`].
-	pub fn handle(&self) -> Arc<Notify> {
-		self.resource.clone()
-	}
-}
-
-/// Wrapper over [`tokio::sync::Notify`] with bounds check.
-#[derive(Debug, Clone)]
-pub struct BoundedSubscriptions {
-	resource: Arc<Notify>,
-	guard: Arc<Semaphore>,
-	max: u32,
-}
-
-impl BoundedSubscriptions {
-	/// Create a new bounded subscription.
-	pub fn new(max_subscriptions: u32) -> Self {
-		Self {
-			resource: Arc::new(Notify::new()),
-			guard: Arc::new(Semaphore::new(max_subscriptions as usize)),
-			max: max_subscriptions,
-		}
-	}
-
-	/// Attempts to acquire a subscription slot.
-	///
-	/// Fails if `max_subscriptions` have been exceeded.
-	pub fn acquire(&self) -> Option<SubscriptionPermit> {
-		Arc::clone(&self.guard)
-			.try_acquire_owned()
-			.ok()
-			.map(|p| SubscriptionPermit { _permit: p, resource: self.resource.clone() })
-	}
-
-	/// Get the maximum number of permitted subscriptions.
-	pub const fn max(&self) -> u32 {
-		self.max
-	}
-
-	/// Close all subscriptions.
-	pub fn close(&self) {
-		self.resource.notify_waiters();
-	}
-}
-
 /// Represent the response to method call.
 #[derive(Debug, Clone)]
 pub struct MethodResponse {
@@ -331,7 +299,6 @@ impl BatchResponse {
 
 #[cfg(test)]
 mod tests {
-	use crate::server::helpers::BoundedSubscriptions;
 
 	use super::{BatchResponseBuilder, BoundedWriter, Id, MethodResponse, Response};
 
@@ -351,20 +318,6 @@ mod tests {
 		assert!(serde_json::to_writer(&mut writer, &"x".repeat(99)).is_err());
 	}
 
-	#[test]
-	fn bounded_subscriptions_work() {
-		let subs = BoundedSubscriptions::new(5);
-		let mut handles = Vec::new();
-
-		for _ in 0..5 {
-			handles.push(subs.acquire().unwrap());
-		}
-
-		assert!(subs.acquire().is_none());
-		handles.swap_remove(0);
-		assert!(subs.acquire().is_some());
-	}
-
 	#[test]
 	fn batch_with_single_works() {
 		let method = MethodResponse::response(Id::Number(1), "a", usize::MAX);

core/src/server/mod.rs

@@ -30,7 +30,5 @@
 pub mod helpers;
 /// Host filtering.
 pub mod host_filtering;
-/// Resource limiting. Create generic "resources" and configure their limits to ensure servers are not overloaded.
-pub mod resource_limiting;
 /// JSON-RPC "modules" group sets of methods that belong together and handles method/subscription registration.
 pub mod rpc_module;