Allow awaiting on server handles

benches/helpers.rs

@@ -70,7 +70,7 @@ pub async fn ws_server(handle: tokio::runtime::Handle) -> (String, jsonrpc_ws_se
 
 /// Run jsonrpsee HTTP server for benchmarks.
 #[cfg(not(feature = "jsonrpc-crate"))]
-pub async fn http_server(handle: tokio::runtime::Handle) -> (String, jsonrpsee::http_server::HttpStopHandle) {
+pub async fn http_server(handle: tokio::runtime::Handle) -> (String, jsonrpsee::http_server::HttpServerHandle) {
 	use jsonrpsee::http_server::{HttpServerBuilder, RpcModule};
 
 	let server = HttpServerBuilder::default()
@@ -88,7 +88,7 @@ pub async fn http_server(handle: tokio::runtime::Handle) -> (String, jsonrpsee::
 
 /// Run jsonrpsee WebSocket server for benchmarks.
 #[cfg(not(feature = "jsonrpc-crate"))]
-pub async fn ws_server(handle: tokio::runtime::Handle) -> (String, jsonrpsee::ws_server::WsStopHandle) {
+pub async fn ws_server(handle: tokio::runtime::Handle) -> (String, jsonrpsee::ws_server::WsServerHandle) {
 	use jsonrpsee::ws_server::{RpcModule, WsServerBuilder};
 
 	let server = WsServerBuilder::default()

examples/http.rs

@@ -26,7 +26,7 @@
 
 use jsonrpsee::{
 	http_client::HttpClientBuilder,
-	http_server::{HttpServerBuilder, HttpStopHandle, RpcModule},
+	http_server::{HttpServerBuilder, HttpServerHandle, RpcModule},
 	rpc_params,
 	types::traits::Client,
 };
@@ -49,12 +49,12 @@ async fn main() -> anyhow::Result<()> {
 	Ok(())
 }
 
-async fn run_server() -> anyhow::Result<(SocketAddr, HttpStopHandle)> {
+async fn run_server() -> anyhow::Result<(SocketAddr, HttpServerHandle)> {
 	let server = HttpServerBuilder::default().build("127.0.0.1:0".parse()?)?;
 	let mut module = RpcModule::new(());
 	module.register_method("say_hello", |_, _| Ok("lo"))?;
 
 	let addr = server.local_addr()?;
-	let stop_handle = server.start(module)?;
-	Ok((addr, stop_handle))
+	let server_handle = server.start(module)?;
+	Ok((addr, server_handle))
 }

examples/proc_macro.rs

@@ -28,7 +28,7 @@ use jsonrpsee::{
 	proc_macros::rpc,
 	types::{async_trait, error::Error, Subscription},
 	ws_client::WsClientBuilder,
-	ws_server::{SubscriptionSink, WsServerBuilder, WsStopHandle},
+	ws_server::{SubscriptionSink, WsServerBuilder, WsServerHandle},
 };
 use std::net::SocketAddr;
 
@@ -89,7 +89,7 @@ async fn main() -> anyhow::Result<()> {
 	Ok(())
 }
 
-async fn run_server() -> anyhow::Result<(SocketAddr, WsStopHandle)> {
+async fn run_server() -> anyhow::Result<(SocketAddr, WsServerHandle)> {
 	let server = WsServerBuilder::default().build("127.0.0.1:0").await?;
 
 	let addr = server.local_addr()?;

http-server/src/lib.rs

@@ -43,7 +43,7 @@ pub use access_control::{
 };
 pub use jsonrpsee_types as types;
 pub use jsonrpsee_utils::server::rpc_module::RpcModule;
-pub use server::{Builder as HttpServerBuilder, Server as HttpServer, StopHandle as HttpStopHandle};
+pub use server::{Builder as HttpServerBuilder, Server as HttpServer, ServerHandle as HttpServerHandle};
 
 #[cfg(test)]
 mod tests;

http-server/src/server.rs

@@ -26,8 +26,7 @@
 
 use crate::{response, AccessControl};
 use futures_channel::mpsc;
-use futures_util::future::join_all;
-use futures_util::stream::StreamExt;
+use futures_util::{future::join_all, stream::StreamExt, FutureExt};
 use hyper::{
 	server::{conn::AddrIncoming, Builder as HyperBuilder},
 	service::{make_service_fn, service_fn},
@@ -49,7 +48,10 @@ use serde_json::value::RawValue;
 use socket2::{Domain, Socket, Type};
 use std::{
 	cmp,
+	future::Future,
 	net::{SocketAddr, TcpListener},
+	pin::Pin,
+	task::{Context, Poll},
 };
 
 /// Builder to create JSON-RPC HTTP server.
@@ -143,24 +145,37 @@ impl Default for Builder {
 	}
 }
 
-/// Handle used to stop the running server.
+/// Handle used to run or stop the server.
 #[derive(Debug)]
-pub struct StopHandle {
+pub struct ServerHandle {
 	stop_sender: mpsc::Sender<()>,
-	stop_handle: Option<tokio::task::JoinHandle<()>>,
+	pub(crate) handle: Option<tokio::task::JoinHandle<()>>,
 }
 
-impl StopHandle {
+impl ServerHandle {
 	/// Requests server to stop. Returns an error if server was already stopped.
 	pub fn stop(mut self) -> Result<tokio::task::JoinHandle<()>, Error> {
-		let stop = self.stop_sender.try_send(()).map(|_| self.stop_handle.take());
+		let stop = self.stop_sender.try_send(()).map(|_| self.handle.take());
 		match stop {
 			Ok(Some(handle)) => Ok(handle),
 			_ => Err(Error::AlreadyStopped),
 		}
 	}
 }
 
+impl Future for ServerHandle {
+	type Output = ();
+
+	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+		let handle = match &mut self.handle {
+			Some(handle) => handle,
+			None => return Poll::Ready(()),
+		};
+
+		handle.poll_unpin(cx).map(|_| ())
+	}
+}
+
 /// An HTTP JSON RPC server.
 #[derive(Debug)]
 pub struct Server {
@@ -185,7 +200,7 @@ impl Server {
 	}
 
 	/// Start the server.
-	pub fn start(mut self, methods: impl Into<Methods>) -> Result<StopHandle, Error> {
+	pub fn start(mut self, methods: impl Into<Methods>) -> Result<ServerHandle, Error> {
 		let max_request_body_size = self.max_request_body_size;
 		let access_control = self.access_control;
 		let (tx, mut rx) = mpsc::channel(1);
@@ -298,7 +313,7 @@ impl Server {
 			let _ = server.with_graceful_shutdown(async move { rx.next().await.map_or((), |_| ()) }).await;
 		});
 
-		Ok(StopHandle { stop_handle: Some(handle), stop_sender: tx })
+		Ok(ServerHandle { handle: Some(handle), stop_sender: tx })
 	}
 }
 

http-server/src/tests.rs

@@ -27,16 +27,17 @@
 #![cfg(test)]
 
 use std::net::SocketAddr;
+use std::time::Duration;
 
 use crate::types::error::{CallError, Error};
-use crate::{server::StopHandle, HttpServerBuilder, RpcModule};
+use crate::{server::ServerHandle, HttpServerBuilder, RpcModule};
 
 use jsonrpsee_test_utils::helpers::*;
 use jsonrpsee_test_utils::mocks::{Id, StatusCode, TestContext};
 use jsonrpsee_test_utils::TimeoutFutureExt;
 use serde_json::Value as JsonValue;
 
-async fn server() -> (SocketAddr, StopHandle) {
+async fn server() -> (SocketAddr, ServerHandle) {
 	let server = HttpServerBuilder::default().build("127.0.0.1:0".parse().unwrap()).unwrap();
 	let ctx = TestContext;
 	let mut module = RpcModule::new(ctx);
@@ -78,8 +79,8 @@ async fn server() -> (SocketAddr, StopHandle) {
 		})
 		.unwrap();
 
-	let stop_handle = server.start(module).unwrap();
-	(addr, stop_handle)
+	let server_handle = server.start(module).unwrap();
+	(addr, server_handle)
 }
 
 #[tokio::test]
@@ -380,6 +381,27 @@ async fn can_register_modules() {
 #[tokio::test]
 async fn stop_works() {
 	let _ = env_logger::try_init();
-	let (_addr, stop_handle) = server().with_default_timeout().await.unwrap();
-	assert!(matches!(stop_handle.stop().unwrap().await, Ok(_)));
+	let (_addr, server_handle) = server().with_default_timeout().await.unwrap();
+	assert!(matches!(server_handle.stop().unwrap().await, Ok(_)));
+}
+
+#[tokio::test]
+async fn run_forever() {
+	const TIMEOUT: Duration = Duration::from_millis(200);
+
+	let _ = env_logger::try_init();
+	let (_addr, server_handle) = server().with_default_timeout().await.unwrap();
+	// Timed out.
+	server_handle.with_timeout(TIMEOUT).await.unwrap_err();
+
+	let (_addr, server_handle) = server().await;
+	server_handle.handle.as_ref().unwrap().abort();
+
+	// Cancelled task is still considered to be finished without errors.
+	// A subject to change.
+	server_handle.with_timeout(TIMEOUT).await.unwrap();
+
+	let (_addr, mut server_handle) = server().with_default_timeout().await.unwrap();
+	server_handle.handle.take();
+	server_handle.with_timeout(TIMEOUT).await.unwrap();
 }

tests/tests/helpers.rs

@@ -25,15 +25,15 @@
 // DEALINGS IN THE SOFTWARE.
 
 use jsonrpsee::{
-	http_server::{HttpServerBuilder, HttpStopHandle},
+	http_server::{HttpServerBuilder, HttpServerHandle},
 	types::Error,
-	ws_server::{WsServerBuilder, WsStopHandle},
+	ws_server::{WsServerBuilder, WsServerHandle},
 	RpcModule,
 };
 use std::net::SocketAddr;
 use std::time::Duration;
 
-pub async fn websocket_server_with_subscription() -> (SocketAddr, WsStopHandle) {
+pub async fn websocket_server_with_subscription() -> (SocketAddr, WsServerHandle) {
 	let server = WsServerBuilder::default().build("127.0.0.1:0").await.unwrap();
 
 	let mut module = RpcModule::new(());
@@ -88,9 +88,9 @@ pub async fn websocket_server_with_subscription() -> (SocketAddr, WsStopHandle)
 		.unwrap();
 
 	let addr = server.local_addr().unwrap();
-	let stop_handle = server.start(module).unwrap();
+	let server_handle = server.start(module).unwrap();
 
-	(addr, stop_handle)
+	(addr, server_handle)
 }
 
 pub async fn websocket_server() -> SocketAddr {
@@ -112,7 +112,7 @@ pub async fn websocket_server() -> SocketAddr {
 	addr
 }
 
-pub async fn http_server() -> (SocketAddr, HttpStopHandle) {
+pub async fn http_server() -> (SocketAddr, HttpServerHandle) {
 	let server = HttpServerBuilder::default().build("127.0.0.1:0".parse().unwrap()).unwrap();
 	let mut module = RpcModule::new(());
 	let addr = server.local_addr().unwrap();

tests/tests/resource_limiting.rs

@@ -26,11 +26,11 @@
 
 use jsonrpsee::{
 	http_client::HttpClientBuilder,
-	http_server::{HttpServerBuilder, HttpStopHandle},
+	http_server::{HttpServerBuilder, HttpServerHandle},
 	proc_macros::rpc,
 	types::{traits::Client, Error},
 	ws_client::WsClientBuilder,
-	ws_server::{WsServerBuilder, WsStopHandle},
+	ws_server::{WsServerBuilder, WsServerHandle},
 	RpcModule,
 };
 use tokio::time::sleep;
@@ -91,7 +91,7 @@ fn module_macro() -> RpcModule<()> {
 	().into_rpc()
 }
 
-async fn websocket_server(module: RpcModule<()>) -> Result<(SocketAddr, WsStopHandle), Error> {
+async fn websocket_server(module: RpcModule<()>) -> Result<(SocketAddr, WsServerHandle), Error> {
 	let server = WsServerBuilder::default()
 		.register_resource("CPU", 6, 2)?
 		.register_resource("MEM", 10, 1)?
@@ -104,7 +104,7 @@ async fn websocket_server(module: RpcModule<()>) -> Result<(SocketAddr, WsStopHa
 	Ok((addr, handle))
 }
 
-async fn http_server(module: RpcModule<()>) -> Result<(SocketAddr, HttpStopHandle), Error> {
+async fn http_server(module: RpcModule<()>) -> Result<(SocketAddr, HttpServerHandle), Error> {
 	let server = HttpServerBuilder::default()
 		.register_resource("CPU", 6, 2)?
 		.register_resource("MEM", 10, 1)?
@@ -128,7 +128,7 @@ fn assert_server_busy(fail: Result<String, Error>) {
 	}
 }
 
-async fn run_tests_on_ws_server(server_addr: SocketAddr, stop_handle: WsStopHandle) {
+async fn run_tests_on_ws_server(server_addr: SocketAddr, server_handle: WsServerHandle) {
 	let server_url = format!("ws://{}", server_addr);
 	let client = WsClientBuilder::default().build(&server_url).await.unwrap();
 
@@ -162,10 +162,10 @@ async fn run_tests_on_ws_server(server_addr: SocketAddr, stop_handle: WsStopHand
 
 	// Client being active prevents the server from shutting down?!
 	drop(client);
-	stop_handle.stop().unwrap().await;
+	server_handle.stop().unwrap().await;
 }
 
-async fn run_tests_on_http_server(server_addr: SocketAddr, stop_handle: HttpStopHandle) {
+async fn run_tests_on_http_server(server_addr: SocketAddr, server_handle: HttpServerHandle) {
 	let server_url = format!("http://{}", server_addr);
 	let client = HttpClientBuilder::default().build(&server_url).unwrap();
 
@@ -190,33 +190,33 @@ async fn run_tests_on_http_server(server_addr: SocketAddr, stop_handle: HttpStop
 
 	assert_eq!(passes, 3);
 
-	stop_handle.stop().unwrap().await.unwrap();
+	server_handle.stop().unwrap().await.unwrap();
 }
 
 #[tokio::test]
 async fn ws_server_with_manual_module() {
-	let (server_addr, stop_handle) = websocket_server(module_manual().unwrap()).await.unwrap();
+	let (server_addr, server_handle) = websocket_server(module_manual().unwrap()).await.unwrap();
 
-	run_tests_on_ws_server(server_addr, stop_handle).await;
+	run_tests_on_ws_server(server_addr, server_handle).await;
 }
 
 #[tokio::test]
 async fn ws_server_with_macro_module() {
-	let (server_addr, stop_handle) = websocket_server(module_macro()).await.unwrap();
+	let (server_addr, server_handle) = websocket_server(module_macro()).await.unwrap();
 
-	run_tests_on_ws_server(server_addr, stop_handle).await;
+	run_tests_on_ws_server(server_addr, server_handle).await;
 }
 
 #[tokio::test]
 async fn http_server_with_manual_module() {
-	let (server_addr, stop_handle) = http_server(module_manual().unwrap()).await.unwrap();
+	let (server_addr, server_handle) = http_server(module_manual().unwrap()).await.unwrap();
 
-	run_tests_on_http_server(server_addr, stop_handle).await;
+	run_tests_on_http_server(server_addr, server_handle).await;
 }
 
 #[tokio::test]
 async fn http_server_with_macro_module() {
-	let (server_addr, stop_handle) = http_server(module_macro()).await.unwrap();
+	let (server_addr, server_handle) = http_server(module_macro()).await.unwrap();
 
-	run_tests_on_http_server(server_addr, stop_handle).await;
+	run_tests_on_http_server(server_addr, server_handle).await;
 }

ws-server/src/future.rs

@@ -166,16 +166,17 @@ impl StopMonitor {
 		self.0.shutdown_requested.load(Ordering::Relaxed)
 	}
 
-	pub(crate) fn handle(&self) -> StopHandle {
-		StopHandle(Arc::downgrade(&self.0))
+	pub(crate) fn handle(&self) -> ServerHandle {
+		ServerHandle(Arc::downgrade(&self.0))
 	}
 }
 
-/// Handle that is able to stop the running server.
+/// Handle that is able to stop the running server or wait for it to finish
+/// its execution.
 #[derive(Debug, Clone)]
-pub struct StopHandle(Weak<MonitorInner>);
+pub struct ServerHandle(Weak<MonitorInner>);
 
-impl StopHandle {
+impl ServerHandle {
 	/// Requests server to stop. Returns an error if server was already stopped.
 	///
 	/// Returns a future that can be awaited for when the server shuts down.
@@ -190,6 +191,16 @@ impl StopHandle {
 	}
 }
 
+impl Future for ServerHandle {
+	type Output = ();
+
+	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+		let mut shutdown_waiter = ShutdownWaiter(self.0.clone());
+
+		shutdown_waiter.poll_unpin(cx)
+	}
+}
+
 /// A `Future` that resolves once the server has stopped.
 #[derive(Debug)]
 pub struct ShutdownWaiter(Weak<MonitorInner>);

ws-server/src/lib.rs

@@ -38,7 +38,7 @@ mod server;
 #[cfg(test)]
 mod tests;
 
-pub use future::{ShutdownWaiter as WsShutdownWaiter, StopHandle as WsStopHandle};
+pub use future::{ServerHandle as WsServerHandle, ShutdownWaiter as WsShutdownWaiter};
 pub use jsonrpsee_types as types;
 pub use jsonrpsee_utils::server::rpc_module::{RpcModule, SubscriptionSink};
 pub use server::{Builder as WsServerBuilder, Server as WsServer};