diff --git a/Cargo.toml b/Cargo.toml index 02cc7746..9d3abc1c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,8 @@ postcard = { workspace = true, features = ["alloc", "use-std"], optional = true noq = { workspace = true, optional = true } # used as a buffer for serialization when using rpc smallvec = { version = "1.14.0", features = ["write"], optional = true } +# used for the Debug impl of rpc::Handler +derive_more = { version = "2", features = ["debug"], optional = true } # used in the test utils to generate noq endpoints rustls = { version = "0.23.5", default-features = false, features = ["std"], optional = true } # used in the test utils to generate noq endpoints @@ -89,7 +91,7 @@ tracing-opentelemetry = "0.33" [features] # enable the remote transport -rpc = ["dep:noq", "dep:postcard", "dep:smallvec", "dep:tracing", "tokio/io-util"] +rpc = ["dep:noq", "dep:postcard", "dep:smallvec", "dep:tracing", "dep:derive_more", "tokio/io-util"] # enable the iroh transport iroh = ["rpc", "dep:iroh", "dep:portable-atomic", "dep:getrandom"] # add helpers to create noq endpoints, for tests and for rpc on localhost diff --git a/examples/0rtt.rs b/examples/0rtt.rs index d155e0c9..9af194d2 100644 --- a/examples/0rtt.rs +++ b/examples/0rtt.rs @@ -240,7 +240,6 @@ mod ping { Client, WithChannels, channel::oneshot, iroh::{Iroh0RttProtocol, IrohProtocol}, - rpc::RemoteService, rpc_requests, }; use serde::{Deserialize, Serialize}; @@ -270,7 +269,7 @@ mod ping { .inner .as_local() .context("can not listen on remote service")?; - Ok(Iroh0RttProtocol::new(EchoProtocol::remote_handler(local))) + Ok(Iroh0RttProtocol::new(local)) } pub fn expose(self) -> Result> { @@ -278,7 +277,7 @@ mod ping { .inner .as_local() .context("can not listen on remote service")?; - Ok(IrohProtocol::new(EchoProtocol::remote_handler(local))) + Ok(IrohProtocol::new(local)) } pub async fn connect( diff --git a/examples/compute.rs b/examples/compute.rs index 15d7da5d..4a12873a 100644 --- a/examples/compute.rs +++ b/examples/compute.rs @@ -9,7 +9,7 @@ use irpc::{ Client, Request, WithChannels, channel::{mpsc, oneshot}, noq::{listen, make_client_endpoint, make_server_endpoint}, - rpc::RemoteService, + rpc::Handler, rpc_requests, }; use n0_future::{ @@ -154,7 +154,7 @@ impl ComputeApi { let Some(local) = self.inner.as_local() else { bail!("cannot listen on a remote service"); }; - let handler = ComputeProtocol::remote_handler(local); + let handler = Handler::from_sender(local); Ok(AbortOnDropHandle::new(task::spawn(listen( endpoint, handler, )))) diff --git a/examples/derive.rs b/examples/derive.rs index 4bc9d27c..3a81e200 100644 --- a/examples/derive.rs +++ b/examples/derive.rs @@ -8,7 +8,6 @@ use irpc::{ Client, WithChannels, channel::{mpsc, oneshot}, noq::{make_client_endpoint, make_server_endpoint}, - rpc::RemoteService, rpc_requests, }; // Import the macro @@ -130,10 +129,7 @@ impl StorageApi { .inner .as_local() .context("cannot listen on remote API")?; - let join_handle = task::spawn(irpc::noq::listen( - endpoint, - StorageProtocol::remote_handler(local), - )); + let join_handle = task::spawn(irpc::noq::listen(endpoint, local)); Ok(AbortOnDropHandle::new(join_handle)) } diff --git a/examples/remote-and-local.rs b/examples/remote-and-local.rs index a1efcac9..cde41a2e 100644 --- a/examples/remote-and-local.rs +++ b/examples/remote-and-local.rs @@ -77,7 +77,6 @@ mod storage { use irpc::{ Client, WithChannels, channel::{mpsc, oneshot}, - rpc::RemoteService, rpc_requests, }; use n0_future::task::AbortOnDropHandle; @@ -183,7 +182,7 @@ mod storage { .client .as_local() .context("can not listen on remote service")?; - Ok(IrohProtocol::new(StorageProtocol::remote_handler(local))) + Ok(IrohProtocol::new(local)) } pub async fn get(&self, key: String) -> irpc::Result> { diff --git a/examples/server-actor.rs b/examples/server-actor.rs index 8d0192d5..7566696a 100644 --- a/examples/server-actor.rs +++ b/examples/server-actor.rs @@ -41,7 +41,7 @@ mod proto { let client = Client::::local(tx); let endpoint = Endpoint::bind(presets::N0).await?; - let protocol = IrohProtocol::with_sender(client.as_local().unwrap()); + let protocol = IrohProtocol::new(client.as_local().unwrap()); let router = Router::builder(endpoint).accept(ALPN, protocol).spawn(); println!("endpoint id: {}", router.endpoint().id()); diff --git a/examples/span_propagation.rs b/examples/span_propagation.rs index 346ef878..68058f7e 100644 --- a/examples/span_propagation.rs +++ b/examples/span_propagation.rs @@ -24,12 +24,12 @@ //! After opening a trace, click "Trace Logs" in the view selector at the top right //! to open a log view that assembles logs from both client and server. -use std::{sync::Arc, time::Duration}; +use std::time::Duration; use anyhow::Result; use clap::{Parser, Subcommand}; use iroh::{Endpoint, EndpointId, endpoint::presets, protocol::Router}; -use irpc::{WithChannels, channel::oneshot, iroh::IrohProtocol, rpc::RemoteService, rpc_requests}; +use irpc::{WithChannels, channel::oneshot, iroh::IrohProtocol, rpc::Handler, rpc_requests}; use opentelemetry::KeyValue; use opentelemetry_otlp::WithExportConfig; use opentelemetry_sdk::{Resource, trace::SdkTracerProvider}; @@ -148,38 +148,35 @@ async fn server(otlp_endpoint: &str) -> Result<()> { // on `rpc_requests` above. let endpoint = Endpoint::bind(presets::N0).await?; let provider = init_tracing("example-server", endpoint.id(), otlp_endpoint)?; - let protocol = IrohProtocol::::new(Arc::new(|req, rx, tx| { - Box::pin(async move { - let msg: Message = ::with_remote_channels(req, rx, tx); - match msg { - Message::Get(msg) => { - // `WithChannels` carries a `span` field that the derive - // macro builds as `info_span!("Get")` with its parent set - // from the propagated remote context, so the handler's - // work shows up under the originating client trace. - let WithChannels { - inner, tx, span, .. - } = msg; - // Hand the response off to a separate task to show that - // `.instrument(span)` keeps work hooked into the same - // trace even after the handler future returns. - tokio::spawn( - async move { - info!(?inner, "handling request"); - // Simulate async work. - tokio::time::sleep(Duration::from_millis(rand::random_range(20..200))) - .await; - let res = inner.0.to_uppercase(); - info!(?res, "generated response"); - tx.send(res).await.ok(); - info!("response sent"); - } - .instrument(span), - ); - } + let protocol = IrohProtocol::::new(Handler::sequential(|msg: Message| async move { + match msg { + Message::Get(msg) => { + // `WithChannels` carries a `span` field that the derive + // macro builds as `info_span!("Get")` with its parent set + // from the propagated remote context, so the handler's + // work shows up under the originating client trace. + let WithChannels { + inner, tx, span, .. + } = msg; + // Hand the response off to a separate task to show that + // `.instrument(span)` keeps work hooked into the same + // trace even after the handler future returns. + tokio::spawn( + async move { + info!(?inner, "handling request"); + // Simulate async work. + tokio::time::sleep(Duration::from_millis(rand::random_range(20..200))) + .await; + let res = inner.0.to_uppercase(); + info!(?res, "generated response"); + tx.send(res).await.ok(); + info!("response sent"); + } + .instrument(span), + ); } - Ok(()) - }) + } + Ok(()) })); let router = Router::builder(endpoint).accept(ALPN, protocol).spawn(); diff --git a/examples/storage.rs b/examples/storage.rs index 7e4e80e4..1b0f6fe7 100644 --- a/examples/storage.rs +++ b/examples/storage.rs @@ -8,7 +8,7 @@ use irpc::{ Channels, Client, Request, Service, WithChannels, channel::{mpsc, none::NoReceiver, oneshot}, noq::{listen, make_client_endpoint, make_server_endpoint}, - rpc::RemoteService, + rpc::{Handler, RemoteService}, }; use n0_future::task::{self, AbortOnDropHandle}; use serde::{Deserialize, Serialize}; @@ -135,7 +135,7 @@ impl StorageApi { let Some(local) = self.inner.as_local() else { bail!("cannot listen on a remote service"); }; - let handler = StorageProtocol::remote_handler(local); + let handler = Handler::from_sender(local); Ok(AbortOnDropHandle::new(task::spawn(listen( endpoint, handler, )))) diff --git a/src/rpc.rs b/src/rpc.rs index 0ae16e72..10e6fe7b 100644 --- a/src/rpc.rs +++ b/src/rpc.rs @@ -449,14 +449,232 @@ impl DynSender for NoqSender { } } -/// Type alias for a handler fn for remote requests -pub type Handler = Arc< - dyn Fn(R, noq::RecvStream, noq::SendStream) -> BoxFuture> +/// The function inside a [`Handler`]. +type HandlerFn = Arc< + dyn Fn(S, noq::RecvStream, noq::SendStream) -> BoxFuture> + Send + Sync + 'static, >; +/// The error that a handler function returns to close the connection. +/// +/// Currently the only constructor is [`Self::close_connection`]. The server +/// then closes the connection. +#[derive(Debug)] +pub struct HandlerError { + inner: HandlerErrorInner, +} + +#[derive(Debug)] +enum HandlerErrorInner { + CloseConnection { code: VarInt, reason: Vec }, +} + +impl std::error::Error for HandlerError {} + +impl std::fmt::Display for HandlerError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match &self.inner { + HandlerErrorInner::CloseConnection { code, reason } => write!( + f, + "handler asked to close the connection with code {code}: {}", + String::from_utf8_lossy(reason) + ), + } + } +} + +impl HandlerError { + /// Creates a [`HandlerError`] that closes the connection with `code` and `reason`. + /// + /// When a handler function returns this error, the server closes the connection + /// immediately and reads no more requests from it. This also ends the streams of + /// all other requests on the connection, including streams that a handler moved + /// into a spawned task. + /// + /// # Examples + /// + /// ``` + /// use irpc::{ + /// channel::oneshot, + /// rpc::{Handler, HandlerError}, + /// rpc_requests, + /// }; + /// use serde::{Deserialize, Serialize}; + /// + /// #[rpc_requests(message = AuthMessage)] + /// #[derive(Debug, Serialize, Deserialize)] + /// enum AuthProtocol { + /// #[rpc(tx = oneshot::Sender<()>)] + /// #[wrap(Auth)] + /// Auth(String), + /// } + /// + /// let handler = Handler::::sequential(|msg| async move { + /// let AuthMessage::Auth(msg) = msg; + /// if msg.inner.0 != "secret" { + /// Err(HandlerError::close_connection(401, "permission denied")) + /// } else { + /// msg.tx.send(()).await.ok(); + /// Ok(()) + /// } + /// }); + /// ``` + pub fn close_connection(code: u32, reason: impl AsRef<[u8]>) -> Self { + Self { + inner: HandlerErrorInner::CloseConnection { + code: VarInt::from_u32(code), + reason: reason.as_ref().to_vec(), + }, + } + } +} + +/// 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::sequential`]: runs a function on each request, one at a time. +/// - [`Handler::raw`]: runs a function on the protocol enum and the streams. +/// +/// A handler function can return `Err(HandlerError)` to close the connection. +/// In all other cases, handle errors internally and return `Ok(())` from +/// the handler function to read the next request from the connection. +/// +/// Often, the protocol type cannot be inferred from the message enum, +/// so you may have to name it: `Handler::::sequential(..)`. +/// +/// # Examples +/// +/// ``` +/// use irpc::{WithChannels, channel::oneshot, rpc::Handler, rpc_requests}; +/// use serde::{Deserialize, Serialize}; +/// +/// #[rpc_requests(message = EchoMessage)] +/// #[derive(Debug, Serialize, Deserialize)] +/// enum EchoProtocol { +/// #[rpc(tx = oneshot::Sender)] +/// #[wrap(Echo)] +/// Echo(String), +/// } +/// +/// let handler: Handler = Handler::sequential(|msg| async move { +/// match msg { +/// EchoMessage::Echo(msg) => { +/// let WithChannels { inner, tx, .. } = msg; +/// tx.send(inner.0).await.ok(); +/// } +/// } +/// Ok(()) +/// }); +/// ``` +#[derive(derive_more::Debug)] +pub struct Handler { + #[debug(skip)] + f: HandlerFn, +} + +impl Clone for Handler { + fn clone(&self) -> Self { + Self { f: self.f.clone() } + } +} + +impl Handler { + /// Creates a handler that runs `f` on the protocol enum and the streams of each request. + /// + /// The requests of a connection run one after the other. If `f` needs the + /// message enum, call [`RemoteService::with_remote_channels`] inside its + /// future: the remote span context is only set while the future runs. + pub fn raw(f: F) -> Self + where + F: Fn(S, noq::RecvStream, noq::SendStream) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let f: HandlerFn = Arc::new(move |msg, rx, tx| Box::pin(f(msg, rx, tx))); + Self { f } + } + + /// Runs this handler on one request. + /// + /// Use it in a [`Handler::raw`] to pass a request to another handler. + pub fn call( + &self, + request: S, + rx: noq::RecvStream, + tx: noq::SendStream, + ) -> impl Future> + Send + 'static { + (self.f)(request, rx, tx) + } + + /// Handles the requests of `connection` until the connection ends. + pub async fn handle_connection( + &self, + connection: &impl IncomingRemoteConnection, + ) -> io::Result<()> { + debug!("connection accepted"); + loop { + let Some((msg, carrier, rx, tx)) = read_request_inner::(connection).await? else { + return Ok(()); + }; + 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(()); + } + } + } + } + } +} + +impl Handler { + /// Creates a handler that sends each request to `local_sender`. + /// + /// A full channel applies backpressure to the connection. + pub fn from_sender(local_sender: impl Into>) -> Self { + let local_sender = local_sender.into(); + Self::sequential(move |msg| { + let local_sender = local_sender.clone(); + async move { + local_sender.send_raw(msg).await.map_err(|err| { + // The receiver is gone. Close as a dropped connection does. + warn!("handler stopped: {err:#}"); + HandlerError::close_connection(0, b"") + }) + } + }) + } + + /// Creates a handler that runs `f` on each request, one at a time per connection. + pub fn sequential(f: F) -> Self + where + F: Fn(S::Message) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let f = Arc::new(f); + Self::raw(move |msg, rx, tx| { + let f = f.clone(); + async move { + // Inside the future: the remote span context is only set while it runs. + f(S::with_remote_channels(msg, rx, tx)).await + } + }) + } +} + +impl From> for Handler { + fn from(local_sender: LocalSender) -> Self { + Self::from_sender(local_sender) + } +} + /// Extension trait to [`Service`] to create a [`Service::Message`] from a [`Service`] /// and a pair of QUIC streams. /// @@ -465,27 +683,13 @@ pub trait RemoteService: Service + Sized { /// Returns the message enum for this request by combining `self` (the protocol enum) /// with a pair of QUIC streams for `tx` and `rx` channels. fn with_remote_channels(self, rx: noq::RecvStream, tx: noq::SendStream) -> Self::Message; - - /// Creates a [`Handler`] that forwards all messages to a [`LocalSender`]. - fn remote_handler(local_sender: LocalSender) -> Handler { - Arc::new(move |msg, rx, tx| { - // `with_remote_channels` reads the task-local span context installed by - // the dispatch loop, so it must run inside the future (which is polled - // within that scope) rather than eagerly here. - let local_sender = local_sender.clone(); - Box::pin(async move { - let msg = Self::with_remote_channels(msg, rx, tx); - local_sender.send_raw(msg).await - }) - }) - } } /// Abstracts over the connections that a server can read requests from. /// /// This is implemented for noq connections, and for iroh connections with and -/// without 0-RTT. It is used by [`read_request`] and [`handle_connection`] to -/// work with all of these. +/// without 0-RTT. It is used by [`read_request`] and [`Handler::handle_connection`] +/// to work with all of these. /// /// This trait is sealed: only irpc can implement it. So irpc can add methods to /// it without a breaking change. @@ -499,23 +703,6 @@ pub trait IncomingRemoteConnection: crate::sealed::Sealed { fn close(&self, error_code: VarInt, reason: &[u8]); } -/// Handles a single connection with the provided `handler`. -/// -/// This function handles requests for a service `S`. The wire format used depends on -/// `S::SPAN_PROPAGATION` - if true, span context is expected in the wire format. -pub async fn handle_connection( - connection: &impl IncomingRemoteConnection, - handler: Handler, -) -> io::Result<()> { - debug!("connection accepted"); - loop { - let Some((msg, carrier, rx, tx)) = read_request_inner::(connection).await? else { - return Ok(()); - }; - crate::span_propagation::scope_remote(carrier, handler(msg, rx, tx)).await?; - } -} - /// Reads a request from a connection and converts it to a message enum. /// /// This combines `read_request_raw` with `RemoteService::with_remote_channels`. @@ -538,7 +725,7 @@ pub async fn read_request( /// This accepts a bi-directional stream from the connection and reads and parses the request. /// /// When `S::SPAN_PROPAGATION` is true, any propagated span context on the wire is -/// silently dropped. Use [`handle_connection`] (or [`read_request`]) if you need +/// silently dropped. Use [`Handler::handle_connection`] (or [`read_request`]) if you need /// the propagated context to reach the generated handler spans. /// /// Returns the parsed request and the stream pair if reading and parsing the request succeeded. diff --git a/src/rpc/iroh_impl.rs b/src/rpc/iroh_impl.rs index 9a896f9d..2fae19c1 100644 --- a/src/rpc/iroh_impl.rs +++ b/src/rpc/iroh_impl.rs @@ -16,8 +16,8 @@ use portable_atomic::{AtomicU64, Ordering}; use tracing::{Instrument, debug, error_span, trace_span, warn}; use crate::{ - LocalSender, RequestError, Service, - rpc::{Handler, IncomingRemoteConnection, RemoteConnection, RemoteService, handle_connection}, + RequestError, Service, + rpc::{Handler, IncomingRemoteConnection, RemoteConnection}, }; impl RemoteConnection for Connection { @@ -152,18 +152,13 @@ impl fmt::Debug for IrohProtocol { } impl IrohProtocol { - pub fn with_sender(local_sender: impl Into>) -> Self - where - S: RemoteService, - { - let handler = S::remote_handler(local_sender.into()); - Self::new(handler) - } - /// Creates a new [`IrohProtocol`] for the `handler`. - pub fn new(handler: Handler) -> Self { + /// + /// `handler` can also be a [`LocalSender`](crate::LocalSender). The protocol then sends + /// each request to it. + pub fn new(handler: impl Into>) -> Self { Self { - handler, + handler: handler.into(), request_id: Default::default(), } } @@ -171,9 +166,11 @@ impl IrohProtocol { impl ProtocolHandler for IrohProtocol { async fn accept(&self, connection: Connection) -> Result<(), AcceptError> { - let handler = self.handler.clone(); let request_id = self.request_id.fetch_add(1, Ordering::AcqRel); - let fut = handle_connection::(&connection, handler).map_err(AcceptError::from_err); + let fut = self + .handler + .handle_connection(&connection) + .map_err(AcceptError::from_err); let remote = connection.remote_id().fmt_short(); let span = trace_span!("rpc", id = request_id, remote = %remote); fut.instrument(span).await @@ -198,18 +195,13 @@ impl fmt::Debug for Iroh0RttProtocol { } impl Iroh0RttProtocol { - pub fn with_sender(local_sender: impl Into>) -> Self - where - S: RemoteService, - { - let handler = S::remote_handler(local_sender.into()); - Self::new(handler) - } - /// Creates a new [`Iroh0RttProtocol`] for the `handler`. - pub fn new(handler: Handler) -> Self { + /// + /// `handler` can also be a [`LocalSender`](crate::LocalSender). The protocol then sends + /// each request to it. + pub fn new(handler: impl Into>) -> Self { Self { - handler, + handler: handler.into(), request_id: Default::default(), } } @@ -218,14 +210,14 @@ impl Iroh0RttProtocol { impl ProtocolHandler for Iroh0RttProtocol { async fn on_accepting(&self, accepting: Accepting) -> Result { let zrtt_conn = accepting.into_0rtt(); - let handler = self.handler.clone(); let request_id = self.request_id.fetch_add(1, Ordering::AcqRel); let span = trace_span!("rpc", id = request_id, remote = tracing::field::Empty); // The remote id of a 0-RTT connection can be unknown. if let Ok(remote) = zrtt_conn.remote_id() { span.record("remote", tracing::field::display(remote.fmt_short())); } - handle_connection::(&zrtt_conn, handler) + self.handler + .handle_connection(&zrtt_conn) .map_err(AcceptError::from_err) .instrument(span) .await?; @@ -269,7 +261,8 @@ impl IncomingRemoteConnection for Connection { /// Utility function to listen for incoming connections and handle them with the provided handler. /// /// The wire format used depends on `S::SPAN_PROPAGATION` - if true, span context is expected. -pub async fn listen(endpoint: iroh::Endpoint, handler: Handler) { +pub async fn listen(endpoint: iroh::Endpoint, handler: impl Into>) { + let handler = handler.into(); let mut request_id = 0u64; let mut tasks = n0_future::task::JoinSet::new(); loop { @@ -291,7 +284,7 @@ pub async fn listen(endpoint: iroh::Endpoint, handler: Handler) { Ok(connection) => { let remote = connection.remote_id().fmt_short(); tracing::Span::current().record("remote", tracing::field::display(remote)); - match handle_connection::(&connection, handler).await { + match handler.handle_connection(&connection).await { Err(err) => warn!("connection closed with error: {err:?}"), Ok(()) => debug!("connection closed"), } diff --git a/src/rpc/noq_impl.rs b/src/rpc/noq_impl.rs index 431ee0c7..62fb45e1 100644 --- a/src/rpc/noq_impl.rs +++ b/src/rpc/noq_impl.rs @@ -7,7 +7,7 @@ use tracing::{Instrument, debug, error_span, warn}; use crate::{ RequestError, Service, - rpc::{Handler, IncomingRemoteConnection, RemoteConnection, handle_connection}, + rpc::{Handler, IncomingRemoteConnection, RemoteConnection}, }; /// A connection to a remote service. @@ -102,7 +102,8 @@ async fn connect_and_open_bi( /// Utility function to listen for incoming connections and handle them with the provided handler. /// /// The wire format used depends on `S::SPAN_PROPAGATION` - if true, span context is expected. -pub async fn listen(endpoint: noq::Endpoint, handler: Handler) { +pub async fn listen(endpoint: noq::Endpoint, handler: impl Into>) { + let handler = handler.into(); let mut request_id = 0u64; let mut tasks = JoinSet::new(); loop { @@ -125,7 +126,7 @@ pub async fn listen(endpoint: noq::Endpoint, handler: Handler) { if let Some(remote) = remote_address(&connection) { tracing::Span::current().record("remote", tracing::field::display(remote)); } - match handle_connection(&connection, handler).await { + match handler.handle_connection(&connection).await { Err(err) => warn!("connection closed with error: {err:?}"), Ok(()) => debug!("connection closed"), } diff --git a/src/span_propagation.rs b/src/span_propagation.rs index 65d320f6..06fe2832 100644 --- a/src/span_propagation.rs +++ b/src/span_propagation.rs @@ -5,7 +5,7 @@ //! but actual OpenTelemetry integration requires the `tracing-opentelemetry` feature. //! //! The propagated context is scoped to a single request handler via a tokio task-local, -//! installed by the dispatch loop in `handle_connection`. This isolates concurrent +//! installed by the dispatch loop in `Handler::handle_connection`. This isolates concurrent //! requests from each other and is robust to thread migration across `.await` points. use std::{collections::HashMap, future::Future}; diff --git a/tests/span_propagation.rs b/tests/span_propagation.rs index 58017b4d..459b1f44 100644 --- a/tests/span_propagation.rs +++ b/tests/span_propagation.rs @@ -1,13 +1,12 @@ #[cfg(feature = "tracing-opentelemetry")] mod span_propagation { - use std::sync::Arc; use iroh::{Endpoint, endpoint::presets, protocol::Router}; use irpc::{ Client, Service, WithChannels, channel::oneshot, iroh::{IrohLazyRemoteConnection, IrohProtocol}, - rpc::RemoteService, + rpc::Handler, rpc_requests, }; use n0_error::{Result, StdResultExt}; @@ -73,29 +72,25 @@ mod span_propagation { async fn listen() -> Result { let endpoint = Endpoint::bind(presets::N0).await?; - let protocol = IrohProtocol::::new(Arc::new(|req, rx, tx| { - Box::pin( - async move { - info!("handle request {req:?}"); - let msg: Message = - ::with_remote_channels(req, rx, tx); - info!("handle request {msg:?}"); - match msg { - Message::Get(msg) => { - let WithChannels { - inner, tx, span, .. - } = msg; - info!("handle request {inner:?} with span {span:?}"); - let _guard = span.enter(); - info!("handle request {inner:?}, entered span"); - tx.send(inner.0.to_uppercase()).await.ok(); - } + let handler = Handler::sequential(|msg: Message| { + async move { + info!("handle request {msg:?}"); + match msg { + Message::Get(msg) => { + let WithChannels { + inner, tx, span, .. + } = msg; + info!("handle request {inner:?} with span {span:?}"); + let _guard = span.enter(); + info!("handle request {inner:?}, entered span"); + tx.send(inner.0.to_uppercase()).await.ok(); } - Ok(()) } - .instrument(info_span!("server-handler")), - ) - })); + Ok(()) + } + .instrument(info_span!("server-handler")) + }); + let protocol = IrohProtocol::::new(handler); let router = Router::builder(endpoint).accept(ALPN, protocol).spawn(); info!("endpoint id: {}", router.endpoint().id()); Ok(router) diff --git a/tests/span_propagation_concurrent.rs b/tests/span_propagation_concurrent.rs index 9badf0f0..c4a75f53 100644 --- a/tests/span_propagation_concurrent.rs +++ b/tests/span_propagation_concurrent.rs @@ -14,14 +14,12 @@ #![cfg(feature = "tracing-opentelemetry")] -use std::sync::Arc; - use iroh::{Endpoint, endpoint::presets, protocol::Router}; use irpc::{ Client, WithChannels, channel::oneshot, iroh::{IrohLazyRemoteConnection, IrohProtocol}, - rpc::RemoteService, + rpc::{Handler, RemoteService}, rpc_requests, }; use n0_error::StdResultExt; @@ -77,21 +75,19 @@ async fn span_propagation_concurrent() -> n0_error::Result<()> { const ALPN: &[u8] = b"test-concurrent"; let endpoint = Endpoint::bind(presets::N0).await?; - let protocol = IrohProtocol::::new(Arc::new(|req, rx, tx| { - Box::pin(async move { - // Yield before constructing the WithChannels so the handler is - // very likely to resume on a different worker thread. - tokio::task::yield_now().await; - let msg: Message = ::with_remote_channels(req, rx, tx); - match msg { - Message::Get(msg) => { - let WithChannels { inner, tx, .. } = msg; - tokio::task::yield_now().await; - tx.send(inner.0.to_uppercase()).await.ok(); - } + let protocol = IrohProtocol::::new(Handler::raw(|req, rx, tx| async move { + // Yield before constructing the WithChannels so the handler is + // very likely to resume on a different worker thread. + tokio::task::yield_now().await; + let msg: Message = ::with_remote_channels(req, rx, tx); + match msg { + Message::Get(msg) => { + let WithChannels { inner, tx, .. } = msg; + tokio::task::yield_now().await; + tx.send(inner.0.to_uppercase()).await.ok(); } - Ok(()) - }) + } + Ok(()) })); let server = Router::builder(endpoint).accept(ALPN, protocol).spawn();