Concurrent polling on async methods

http-server/src/server.rs

@@ -26,6 +26,7 @@
 
 use crate::{response, AccessControl};
 use futures_channel::mpsc;
+use futures_util::future::join_all;
 use futures_util::{lock::Mutex, stream::StreamExt, SinkExt};
 use hyper::{
 	server::{conn::AddrIncoming, Builder as HyperBuilder},
@@ -214,13 +215,16 @@ impl Server {
 
 						// NOTE(niklasad1): it's a channel because it's needed for batch requests.
 						let (tx, mut rx) = mpsc::unbounded::<String>();
+
 						type Notif<'a> = JsonRpcNotification<'a, Option<&'a RawValue>>;
 
 						// Single request or notification
 						if is_single {
 							if let Ok(req) = serde_json::from_slice::<JsonRpcRequest>(&body) {
 								// NOTE: we don't need to track connection id on HTTP, so using hardcoded 0 here.
-								methods.execute(&tx, req, 0).await;
+								if let Some(fut) = methods.execute(&tx, req, 0) {
+									fut.await;
+								}
 							} else if let Ok(_req) = serde_json::from_slice::<Notif>(&body) {
 								return Ok::<_, HyperError>(response::ok_response("".into()));
 							} else {
@@ -231,9 +235,7 @@ impl Server {
 						// Batch of requests or notifications
 						} else if let Ok(batch) = serde_json::from_slice::<Vec<JsonRpcRequest>>(&body) {
 							if !batch.is_empty() {
-								for req in batch {
-									methods.execute(&tx, req, 0).await;
-								}
+								join_all(batch.into_iter().filter_map(|req| methods.execute(&tx, req, 0))).await;
 							} else {
 								// "If the batch rpc call itself fails to be recognized as an valid JSON or as an
 								// Array with at least one value, the response from the Server MUST be a single

utils/Cargo.toml

@@ -9,6 +9,7 @@ license = "MIT"
 [dependencies]
 beef = { version = "0.5.1", features = ["impl_serde"] }
 thiserror = { version = "1", optional = true }
+tokio = { version = "1", features = ["macros"], optional = true }
 futures-channel = { version = "0.3.14", default-features = false, optional = true }
 futures-util = { version = "0.3.14", default-features = false, optional = true }
 hyper = { version = "0.14.10", default-features = false, features = ["stream"], optional = true }
@@ -33,9 +34,9 @@ server = [
 	"serde_json",
 	"log",
 	"parking_lot",
-	"rand"
+	"rand",
+	"tokio"
 ]
 
 [dev-dependencies]
 serde_json = "1.0"
-tokio = { version = "1", features = ["macros"] }

utils/src/server/rpc_module.rs

@@ -21,13 +21,10 @@ use std::sync::Arc;
 /// implemented as a function pointer to a `Fn` function taking four arguments:
 /// the `id`, `params`, a channel the function uses to communicate the result (or error)
 /// back to `jsonrpsee`, and the connection ID (useful for the websocket transport).
-pub type SyncMethod = Arc<dyn Send + Sync + Fn(Id, RpcParams, &MethodSink, ConnectionId) -> Result<(), Error>>;
+pub type SyncMethod = Arc<dyn Send + Sync + Fn(Id, RpcParams, &MethodSink, ConnectionId)>;
 /// Similar to [`SyncMethod`], but represents an asynchronous handler.
-pub type AsyncMethod = Arc<
-	dyn Send
-		+ Sync
-		+ Fn(Id<'static>, RpcParams<'static>, MethodSink, ConnectionId) -> BoxFuture<'static, Result<(), Error>>,
->;
+pub type AsyncMethod<'a> =
+	Arc<dyn Send + Sync + Fn(Id<'a>, RpcParams<'a>, MethodSink, ConnectionId) -> BoxFuture<'a, ()>>;
 /// Connection ID, used for stateful protocol such as WebSockets.
 /// For stateless protocols such as http it's unused, so feel free to set it some hardcoded value.
 pub type ConnectionId = usize;
@@ -51,29 +48,33 @@ pub enum MethodCallback {
 	/// Synchronous method handler.
 	Sync(SyncMethod),
 	/// Asynchronous method handler.
-	Async(AsyncMethod),
+	Async(AsyncMethod<'static>),
 }
 
 impl MethodCallback {
 	/// Execute the callback, sending the resulting JSON (success or error) to the specified sink.
-	pub async fn execute(&self, tx: &MethodSink, req: JsonRpcRequest<'_>, conn_id: ConnectionId) {
+	pub fn execute(
+		&self,
+		tx: &MethodSink,
+		req: JsonRpcRequest<'_>,
+		conn_id: ConnectionId,
+	) -> Option<BoxFuture<'static, ()>> {
 		let id = req.id.clone();
 		let params = RpcParams::new(req.params.map(|params| params.get()));
 
-		let result = match self {
-			MethodCallback::Sync(callback) => (callback)(req.id.clone(), params, tx, conn_id),
+		match self {
+			MethodCallback::Sync(callback) => {
+				(callback)(id, params, tx, conn_id);
+
+				None
+			}
 			MethodCallback::Async(callback) => {
 				let tx = tx.clone();
 				let params = params.into_owned();
-				let id = req.id.into_owned();
+				let id = id.into_owned();
 
-				(callback)(id, params, tx, conn_id).await
+				Some((callback)(id, params, tx, conn_id))
 			}
-		};
-
-		if let Err(err) = result {
-			log::error!("execution of method call '{}' failed: {:?}, request id={:?}", req.method, err, id);
-			send_error(id, tx, JsonRpcErrorCode::ServerError(-1).into());
 		}
 	}
 }
@@ -136,10 +137,18 @@ impl Methods {
 	}
 
 	/// Attempt to execute a callback, sending the resulting JSON (success or error) to the specified sink.
-	pub async fn execute(&self, tx: &MethodSink, req: JsonRpcRequest<'_>, conn_id: ConnectionId) {
+	pub fn execute(
+		&self,
+		tx: &MethodSink,
+		req: JsonRpcRequest<'_>,
+		conn_id: ConnectionId,
+	) -> Option<BoxFuture<'static, ()>> {
 		match self.callbacks.get(&*req.method) {
-			Some(callback) => callback.execute(tx, req, conn_id).await,
-			None => send_error(req.id, tx, JsonRpcErrorCode::MethodNotFound.into()),
+			Some(callback) => callback.execute(tx, req, conn_id),
+			None => {
+				send_error(req.id, tx, JsonRpcErrorCode::MethodNotFound.into());
+				None
+			}
 		}
 	}
 
@@ -155,7 +164,9 @@ impl Methods {
 
 		let (tx, mut rx) = mpsc::unbounded();
 
-		self.execute(&tx, req, 0).await;
+		if let Some(fut) = self.execute(&tx, req, 0) {
+			fut.await;
+		}
 
 		rx.next().await
 	}
@@ -233,8 +244,6 @@ impl<Context: Send + Sync + 'static> RpcModule<Context> {
 						send_error(id, tx, err)
 					}
 				};
-
-				Ok(())
 			})),
 		);
 
@@ -277,7 +286,6 @@ impl<Context: Send + Sync + 'static> RpcModule<Context> {
 							send_error(id, &tx, err)
 						}
 					};
-					Ok(())
 				};
 				future.boxed()
 			})),
@@ -344,15 +352,23 @@ impl<Context: Send + Sync + 'static> RpcModule<Context> {
 						sub_id
 					};
 
-					send_response(id, method_sink, sub_id);
+					send_response(id.clone(), method_sink, sub_id);
 					let sink = SubscriptionSink {
 						inner: method_sink.clone(),
 						method: subscribe_method_name,
 						subscribers: subscribers.clone(),
 						uniq_sub: SubscriptionKey { conn_id, sub_id },
 						is_connected: Some(conn_tx),
 					};
-					callback(params, sink, ctx.clone())
+					if let Err(err) = callback(params, sink, ctx.clone()) {
+						log::error!(
+							"subscribe call '{}' failed: {:?}, request id={:?}",
+							subscribe_method_name,
+							err,
+							id
+						);
+						send_error(id, method_sink, JsonRpcErrorCode::ServerError(-1).into());
+					}
 				})),
 			);
 		}
@@ -361,11 +377,20 @@ impl<Context: Send + Sync + 'static> RpcModule<Context> {
 			self.methods.mut_callbacks().insert(
 				unsubscribe_method_name,
 				MethodCallback::Sync(Arc::new(move |id, params, tx, conn_id| {
-					let sub_id = params.one()?;
+					let sub_id = match params.one() {
+						Ok(sub_id) => sub_id,
+						Err(_) => {
+							log::error!(
+								"unsubscribe call '{}' failed: couldn't parse subscription id, request id={:?}",
+								unsubscribe_method_name,
+								id
+							);
+							send_error(id, tx, JsonRpcErrorCode::ServerError(-1).into());
+							return;
+						}
+					};
 					subscribers.lock().remove(&SubscriptionKey { conn_id, sub_id });
 					send_response(id, tx, "Unsubscribed");
-
-					Ok(())
 				})),
 			);
 		}

ws-server/src/future.rs

@@ -0,0 +1,132 @@
+// Copyright 2019 Parity Technologies (UK) Ltd.
+//
+// Permission is hereby granted, free of charge, to any
+// person obtaining a copy of this software and associated
+// documentation files (the "Software"), to deal in the
+// Software without restriction, including without
+// limitation the rights to use, copy, modify, merge,
+// publish, distribute, sublicense, and/or sell copies of
+// the Software, and to permit persons to whom the Software
+// is furnished to do so, subject to the following
+// conditions:
+//
+// The above copyright notice and this permission notice
+// shall be included in all copies or substantial portions
+// of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
+// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
+// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+//! Utilities for handling async code.
+
+use futures_util::future::FutureExt;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+/// This is a flexible collection of futures that need to be driven to completion
+/// alongside some other future, such as connection handlers that need to be
+/// handled along with a listener for new connections.
+///
+/// In order to `.await` on these futures and drive them to completion, call
+/// `select_with` providing some other future, the result of which you need.
+pub(crate) struct FutureDriver<F> {
+	futures: Vec<F>,
+}
+
+impl<F> Default for FutureDriver<F> {
+	fn default() -> Self {
+		FutureDriver { futures: Vec::new() }
+	}
+}
+
+impl<F> FutureDriver<F> {
+	/// Get the count of remaining futures on this driver
+	pub(crate) fn count(&self) -> usize {
+		self.futures.len()
+	}
+
+	/// Add a new future to this driver
+	pub(crate) fn add(&mut self, future: F) {
+		self.futures.push(future);
+	}
+}
+
+impl<F> FutureDriver<F>
+where
+	F: Future + Unpin,
+{
+	pub(crate) async fn select_with<S: Future>(&mut self, selector: S) -> S::Output {
+		tokio::pin!(selector);
+
+		DriverSelect { selector, driver: self }.await
+	}
+
+	fn drive(&mut self, cx: &mut Context) {
+		let mut i = 0;
+
+		while i < self.futures.len() {
+			if self.futures[i].poll_unpin(cx).is_ready() {
+				// Using `swap_remove` since we don't care about ordering
+				// but we do care about removing being `O(1)`.
+				//
+				// We don't increment `i` in this branch, since we now
+				// have a shorter length, and potentially a new value at
+				// current index
+				self.futures.swap_remove(i);
+			} else {
+				i += 1;
+			}
+		}
+	}
+}
+
+impl<F> Future for FutureDriver<F>
+where
+	F: Future + Unpin,
+{
+	type Output = ();
+
+	fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
+		let this = Pin::into_inner(self);
+
+		this.drive(cx);
+
+		if this.futures.len() == 0 {
+			Poll::Ready(())
+		} else {
+			Poll::Pending
+		}
+	}
+}
+
+/// This is a glorified select `Future` that will attempt to drive all
+/// connection futures `F` to completion on each `poll`, while also
+/// handling incoming connections.
+struct DriverSelect<'a, S, F> {
+	selector: S,
+	driver: &'a mut FutureDriver<F>,
+}
+
+impl<'a, R, F> Future for DriverSelect<'a, R, F>
+where
+	R: Future + Unpin,
+	F: Future + Unpin,
+{
+	type Output = R::Output;
+
+	fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
+		let this = Pin::into_inner(self);
+
+		this.driver.drive(cx);
+
+		this.selector.poll_unpin(cx)
+	}
+}

ws-server/src/lib.rs

@@ -32,6 +32,7 @@
 
 extern crate alloc;
 
+mod future;
 mod server;
 
 #[cfg(test)]