diff --git a/src/rpc.rs b/src/rpc.rs index 10e6fe7..35cf33f 100644 --- a/src/rpc.rs +++ b/src/rpc.rs @@ -4,11 +4,17 @@ //! streams. The [`crate::noq`] module has the transport with dial by socket //! address. The [`crate::iroh`] module has the transport with dial by endpoint id. use std::{ - fmt::Debug, future::Future, io, marker::PhantomData, ops::DerefMut, pin::Pin, sync::Arc, + fmt::Debug, + future::Future, + io, + marker::PhantomData, + ops::DerefMut, + pin::{Pin, pin}, + sync::Arc, }; use n0_error::{e, stack_error}; -use n0_future::future::Boxed as BoxFuture; +use n0_future::{FuturesUnorderedBounded, StreamExt, future::Boxed as BoxFuture}; use noq::{ConnectionError, VarInt}; use serde::de::DeserializeOwned; use smallvec::SmallVec; @@ -531,11 +537,22 @@ impl HandlerError { } } +/// How the server loop runs the requests of one connection. +#[derive(Debug, Clone, Copy)] +enum Mode { + /// Runs each request in the loop, before the loop reads the next request. + Sequential, + /// Runs at most this many requests of a connection at the same time. + Concurrent(usize), +} + /// Handles the requests that a server reads from a connection. /// /// A server uses one handler for all its connections. Create one with: /// /// - [`Handler::from_sender`]: sends each request to a [`LocalSender`]. +/// - [`Handler::concurrent`]: runs a function on each request, with a limit +/// of concurrent requests per connection. /// - [`Handler::sequential`]: runs a function on each request, one at a time. /// - [`Handler::raw`]: runs a function on the protocol enum and the streams. /// @@ -560,7 +577,7 @@ impl HandlerError { /// Echo(String), /// } /// -/// let handler: Handler = Handler::sequential(|msg| async move { +/// let handler: Handler = Handler::concurrent(16, |msg| async move { /// match msg { /// EchoMessage::Echo(msg) => { /// let WithChannels { inner, tx, .. } = msg; @@ -574,11 +591,15 @@ impl HandlerError { pub struct Handler { #[debug(skip)] f: HandlerFn, + mode: Mode, } impl Clone for Handler { fn clone(&self) -> Self { - Self { f: self.f.clone() } + Self { + f: self.f.clone(), + mode: self.mode, + } } } @@ -594,12 +615,17 @@ impl Handler { Fut: Future> + Send + 'static, { let f: HandlerFn = Arc::new(move |msg, rx, tx| Box::pin(f(msg, rx, tx))); - Self { f } + Self { + f, + mode: Mode::Sequential, + } } /// Runs this handler on one request. /// - /// Use it in a [`Handler::raw`] to pass a request to another handler. + /// Use it in a [`Handler::raw`] to pass a request to another handler. The + /// request runs in the current task, so the limit of a + /// [`Handler::concurrent`] does not apply. pub fn call( &self, request: S, @@ -610,11 +636,23 @@ impl Handler { } /// Handles the requests of `connection` until the connection ends. + /// + /// When the connection ends, this waits for the requests that still run. pub async fn handle_connection( &self, connection: &impl IncomingRemoteConnection, ) -> io::Result<()> { debug!("connection accepted"); + match self.mode { + Mode::Sequential => self.run_sequential(connection).await, + Mode::Concurrent(max_concurrent) => { + self.run_concurrent(connection, max_concurrent).await + } + } + } + + /// Runs the requests of `connection` one after the other. + async fn run_sequential(&self, connection: &impl IncomingRemoteConnection) -> io::Result<()> { loop { let Some((msg, carrier, rx, tx)) = read_request_inner::(connection).await? else { return Ok(()); @@ -622,15 +660,51 @@ impl Handler { let fut = (self.f)(msg, rx, tx); let fut = crate::span_propagation::scope_remote(carrier, fut); if let Err(err) = fut.await { - match err.inner { - HandlerErrorInner::CloseConnection { code, reason } => { - debug!(%code, "handler closed the connection"); - connection.close(code, &reason); - return Ok(()); + close_connection(connection, err); + return Ok(()); + } + } + } + + /// Runs the requests of `connection` at the same time, at most `max_concurrent` of them. + async fn run_concurrent( + &self, + connection: &impl IncomingRemoteConnection, + max_concurrent: usize, + ) -> io::Result<()> { + let mut futs = FuturesUnorderedBounded::new(max_concurrent); + // Kept across iterations, so that a partly read request is not cancelled. + let mut read = pin!(read_request_inner::(connection)); + let res = loop { + tokio::select! { + biased; + Some(res) = futs.next(), if !futs.is_empty() => { + if let Err(err) = res { + close_connection(connection, err); + break Ok(()); + } + } + next = &mut read, if futs.len() < max_concurrent => { + read.set(read_request_inner::(connection)); + match next { + Err(err) => break Err(err), + Ok(None) => break Ok(()), + Ok(Some((msg, carrier, rx, tx))) => { + let fut = crate::span_propagation::scope_remote(carrier, (self.f)(msg, rx, tx)); + futs.push(fut); + } } } } + }; + // Complete the running requests, also after a graceful close: the remote + // can close right after a `Client::notify`. + while let Some(res) = futs.next().await { + if let Err(err) = res { + close_connection(connection, err); + } } + res } } @@ -652,6 +726,27 @@ impl Handler { }) } + /// Creates a handler that runs `f` on each request, at most `max_concurrent` at a time. + /// + /// The limit applies per connection. The requests run in the task of the + /// connection, so a request that blocks delays the others. Spawn a task in + /// `f` for such work. + /// + /// # Panics + /// + /// Panics if `max_concurrent` is zero. + pub fn concurrent(max_concurrent: usize, f: F) -> Self + where + F: Fn(S::Message) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + assert!(max_concurrent > 0, "max_concurrent must not be zero"); + Self { + mode: Mode::Concurrent(max_concurrent), + ..Self::sequential(f) + } + } + /// Creates a handler that runs `f` on each request, one at a time per connection. pub fn sequential(f: F) -> Self where @@ -703,6 +798,16 @@ pub trait IncomingRemoteConnection: crate::sealed::Sealed { fn close(&self, error_code: VarInt, reason: &[u8]); } +/// Closes `connection` as the error of a handler asks. +fn close_connection(connection: &impl IncomingRemoteConnection, err: HandlerError) { + match err.inner { + HandlerErrorInner::CloseConnection { code, reason } => { + debug!(%code, "handler closed the connection"); + connection.close(code, &reason); + } + } +} + /// Reads a request from a connection and converts it to a message enum. /// /// This combines `read_request_raw` with `RemoteService::with_remote_channels`. diff --git a/tests/handler.rs b/tests/handler.rs new file mode 100644 index 0000000..697e6aa --- /dev/null +++ b/tests/handler.rs @@ -0,0 +1,131 @@ +#![cfg(all(feature = "noq_endpoint_setup", feature = "derive", feature = "spans"))] + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use irpc::{Client, WithChannels, channel::oneshot, noq::listen, rpc::Handler, rpc_requests}; +use n0_future::task::AbortOnDropHandle; +use serde::{Deserialize, Serialize}; +use testresult::TestResult; +use tokio::sync::watch; + +#[allow(dead_code, reason = "this test uses only some of the helpers")] +mod common; +use common::*; + +#[rpc_requests(message = BlockMessage)] +#[derive(Debug, Serialize, Deserialize)] +enum BlockProtocol { + /// Waits until the test opens the gate, then replies. + #[rpc(tx = oneshot::Sender<()>)] + #[wrap(Block)] + Block, +} + +/// Blocks requests until the test opens it, and counts the requests that wait. +#[derive(Debug)] +struct Gate { + waiting: watch::Sender, + max_waiting: AtomicUsize, + passed: AtomicUsize, + open: watch::Sender, +} + +impl Gate { + fn new() -> Arc { + Arc::new(Self { + waiting: watch::Sender::new(0), + max_waiting: AtomicUsize::new(0), + passed: AtomicUsize::new(0), + open: watch::Sender::new(false), + }) + } + + /// Returns a concurrent handler that blocks each request at this gate. + fn handler(self: &Arc, max_concurrent: usize) -> Handler { + let gate = self.clone(); + Handler::concurrent(max_concurrent, move |msg| { + let gate = gate.clone(); + async move { + let BlockMessage::Block(WithChannels { tx, .. }) = msg; + gate.waiting.send_modify(|n| { + *n += 1; + gate.max_waiting.fetch_max(*n, Ordering::SeqCst); + }); + gate.open.subscribe().wait_for(|open| *open).await.ok(); + gate.waiting.send_modify(|n| *n -= 1); + gate.passed.fetch_add(1, Ordering::SeqCst); + tx.send(()).await.ok(); + Ok(()) + } + }) + } + + /// Waits until `n` requests wait at the gate. + async fn wait_for(&self, n: usize) -> TestResult<()> { + let mut waiting = self.waiting.subscribe(); + tokio::time::timeout(Duration::from_secs(10), waiting.wait_for(|w| *w == n)).await??; + Ok(()) + } +} + +/// Sends `requests` requests at once to a handler with `max_concurrent`. +/// +/// Returns how many of them ran at the same time. All requests use one connection. +async fn max_running(max_concurrent: usize, requests: usize) -> TestResult { + let (server, client, server_addr) = create_connected_endpoints()?; + let gate = Gate::new(); + let _server = + AbortOnDropHandle::new(tokio::spawn(listen(server, gate.handler(max_concurrent)))); + let conn = client.connect(server_addr, "localhost")?.await?; + let client = Client::::boxed(conn); + let calls = (0..requests).map(|_| client.rpc(Block)); + let calls = tokio::spawn(futures_util::future::join_all(calls)); + gate.wait_for(max_concurrent.min(requests)).await?; + gate.open.send_replace(true); + for res in calls.await? { + res?; + } + Ok(gate.max_waiting.load(Ordering::SeqCst)) +} + +#[tokio::test] +async fn concurrent_handler_runs_requests_up_to_the_limit() -> TestResult<()> { + assert_eq!(max_running(3, 8).await?, 3); + assert_eq!(max_running(64, 8).await?, 8); + Ok(()) +} + +#[tokio::test] +async fn concurrent_handler_completes_requests_after_the_connection_closes() -> TestResult<()> { + let (server, client, server_addr) = create_connected_endpoints()?; + let gate = Gate::new(); + let handler = gate.handler(4); + let (closed_tx, closed_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let conn = server.accept().await.expect("incoming connection").await?; + let closed = conn.clone(); + tokio::spawn(async move { + closed.closed().await; + closed_tx.send(()).ok(); + }); + handler.handle_connection(&conn).await + }); + let conn = client.connect(server_addr, "localhost")?.await?; + let client = Client::::boxed(conn.clone()); + let _call = tokio::spawn(async move { client.rpc(Block).await }); + gate.wait_for(1).await?; + conn.close(0u32.into(), b""); + // Open the gate only after the server sees the close, so that the request + // still runs when the connection ends. + closed_rx.await?; + gate.open.send_replace(true); + server.await??; + assert_eq!(gate.passed.load(Ordering::SeqCst), 1); + Ok(()) +}