add propagate-span attribute to rpc_request

rread · rread · commit e325e1d0e807 · 2026-01-15T18:07:46.000-05:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -46,7 +46,7 @@ n0-future = { workspace = true }
 futures-util = { workspace = true, optional = true }
 # for the derive reexport/feature
 irpc-derive = { version = "0.8.0", path = "./irpc-derive", optional = true }
-# for remote span propagation when both spans and rpc are enabled
+# for remote span propagation when use-tracing-opentelemetry feature is enabled
 opentelemetry = { version = "0.31", optional = true }
 tracing-opentelemetry = { version = "0.32", optional = true }
 
@@ -67,14 +67,16 @@ testresult = "0.4.1"
 
 [features]
 # enable the remote transport
-rpc = ["dep:quinn", "dep:postcard", "dep:anyhow", "dep:smallvec", "dep:tracing", "tokio/io-util", "dep:opentelemetry", "dep:tracing-opentelemetry", "irpc-derive?/rpc"]
+rpc = ["dep:quinn", "dep:postcard", "dep:anyhow", "dep:smallvec", "dep:tracing", "tokio/io-util"]
 # add test utilities
 quinn_endpoint_setup = ["rpc", "dep:rustls", "dep:rcgen", "dep:anyhow", "dep:futures-buffered", "quinn/rustls-ring"]
 # pick up parent span when creating channel messages
-spans = ["dep:tracing", "irpc-derive?/spans"]
+spans = ["dep:tracing"]
 stream = ["dep:futures-util"]
 derive = ["dep:irpc-derive"]
 varint-util = ["dep:postcard", "dep:smallvec", "tokio/io-util"]
+# enable OpenTelemetry span context propagation across remote connections
+use-tracing-opentelemetry = ["dep:opentelemetry", "dep:tracing-opentelemetry", "rpc"]
 default = ["rpc", "quinn_endpoint_setup", "spans", "stream", "derive"]
 
 [[example]]
diff --git a/irpc-derive/Cargo.toml b/irpc-derive/Cargo.toml
@@ -18,7 +18,4 @@ quote = "1"
 proc-macro2 = "1"
 
 [features]
-# These features should match irpc's features for proper code generation
-spans = []
-rpc = []
-default = ["spans", "rpc"]
+default = []
diff --git a/irpc-derive/src/lib.rs b/irpc-derive/src/lib.rs
@@ -152,15 +152,21 @@ pub fn rpc_requests(attr: TokenStream, item: TokenStream) -> TokenStream {
         let message_from_impls =
             generate_message_enum_from_impls(message_enum_name, &variants_with_attr, enum_name);
 
+        let span_propagation = args.span_propagation;
         let service_impl = quote! {
             impl ::irpc::Service for #enum_name {
                 type Message = #message_enum_name;
+                const SPAN_PROPAGATION: bool = #span_propagation;
             }
         };
 
         let remote_service_impl = if !args.no_rpc {
-            let block =
-                generate_remote_service_impl(message_enum_name, enum_name, &variants_with_attr);
+            let block = generate_remote_service_impl(
+                message_enum_name,
+                enum_name,
+                &variants_with_attr,
+                args.span_propagation,
+            );
             quote! {
                 #cfg_feature_rpc
                 #block
@@ -280,39 +286,39 @@ fn generate_message_enum_from_impls(
 }
 
 /// Generate `RemoteService` impl for message enums.
+///
+/// When `span_propagation` is true, the generated code will create spans for each
+/// request and set their parent from the propagated remote context.
 fn generate_remote_service_impl(
     message_enum_name: &Ident,
     proto_enum_name: &Ident,
     variants_with_attr: &[(Ident, Type)],
+    span_propagation: bool,
 ) -> TokenStream2 {
-    // Generate match arms that set the span parent for each variant
-    #[cfg(all(feature = "spans", feature = "rpc"))]
+    // Generate match arms for each variant
     let variants = variants_with_attr
         .iter()
         .map(|(variant_name, _inner_type)| {
             let span_name = variant_name.to_string();
-            quote! {
-                #proto_enum_name::#variant_name(msg) => {
-                    // Create a span for this specific RPC operation
-                    let span = ::tracing::info_span!(#span_name);
-                    // Set its parent to the propagated remote context if available
-                    if let Some(ctx) = ::irpc::span_propagation::take_remote_span_context() {
-                        use ::irpc::span_propagation::OpenTelemetrySpanExt;
-                        let _ = span.set_parent(ctx);
+
+            if span_propagation {
+                // When span_propagation is enabled, create spans and set parent from remote context
+                quote! {
+                    #proto_enum_name::#variant_name(msg) => {
+                        // Create a span for this specific RPC operation
+                        let span = ::tracing::info_span!(#span_name);
+                        // Set its parent to the propagated remote context if available
+                        ::irpc::span_propagation::set_span_parent_from_remote(&span);
+                        let _guard = span.enter();
+                        #message_enum_name::from(::irpc::WithChannels::from((msg, tx, rx)))
                     }
-                    let _guard = span.enter();
-                    #message_enum_name::from(::irpc::WithChannels::from((msg, tx, rx)))
                 }
-            }
-        });
-
-    #[cfg(not(all(feature = "spans", feature = "rpc")))]
-    let variants = variants_with_attr
-        .iter()
-        .map(|(variant_name, _inner_type)| {
-            quote! {
-                #proto_enum_name::#variant_name(msg) => {
-                    #message_enum_name::from(::irpc::WithChannels::from((msg, tx, rx)))
+            } else {
+                // When span_propagation is disabled, just create the message
+                quote! {
+                    #proto_enum_name::#variant_name(msg) => {
+                        #message_enum_name::from(::irpc::WithChannels::from((msg, tx, rx)))
+                    }
                 }
             }
         });
@@ -361,6 +367,8 @@ struct MacroArgs {
     rpc_feature: Option<String>,
     no_rpc: bool,
     no_spans: bool,
+    /// When true, includes span context in the wire format and enables span propagation.
+    span_propagation: bool,
 }
 
 impl Parse for MacroArgs {
@@ -396,6 +404,9 @@ impl Parse for MacroArgs {
                 "no_spans" => {
                     this.no_spans = true;
                 }
+                "span_propagation" => {
+                    this.span_propagation = true;
+                }
                 _ => {
                     return syn_err(arg.span(), format!("Unknown parameter: {arg}"));
                 }
diff --git a/irpc-iroh/src/lib.rs b/irpc-iroh/src/lib.rs
@@ -14,10 +14,9 @@ use irpc::{
         MAX_MESSAGE_SIZE,
     },
     util::AsyncReadVarintExt,
-    LocalSender, RequestError,
+    LocalSender, RequestError, Service,
 };
 use n0_future::{future::Boxed as BoxFuture, TryFutureExt};
-use serde::de::DeserializeOwned;
 use tracing::{debug, error_span, trace, trace_span, warn, Instrument};
 
 /// Returns a client that connects to a irpc service using an [`iroh::Endpoint`].
@@ -102,8 +101,8 @@ async fn connect_and_open_bi(
 /// A [`ProtocolHandler`] for an irpc protocol.
 ///
 /// Can be added to an [`iroh::protocol::Router`] to handle incoming connections for an ALPN string.
-pub struct IrohProtocol<R> {
-    handler: Handler<R>,
+pub struct IrohProtocol<S> {
+    handler: Handler<S>,
     request_id: AtomicU64,
 }
 
@@ -113,25 +112,25 @@ impl<T> fmt::Debug for IrohProtocol<T> {
     }
 }
 
-impl<R: DeserializeOwned + Send + 'static> IrohProtocol<R> {
-    pub fn with_sender(local_sender: impl Into<LocalSender<R>>) -> Self
+impl<S: Service> IrohProtocol<S> {
+    pub fn with_sender(local_sender: impl Into<LocalSender<S>>) -> Self
     where
-        R: RemoteService,
+        S: RemoteService,
     {
-        let handler = R::remote_handler(local_sender.into());
+        let handler = S::remote_handler(local_sender.into());
         Self::new(handler)
     }
 
     /// Creates a new [`IrohProtocol`] for the `handler`.
-    pub fn new(handler: Handler<R>) -> Self {
+    pub fn new(handler: Handler<S>) -> Self {
         Self {
             handler,
             request_id: Default::default(),
         }
     }
 }
 
-impl<R: DeserializeOwned + Send + 'static> ProtocolHandler for IrohProtocol<R> {
+impl<S: Service> ProtocolHandler for IrohProtocol<S> {
     fn accept(
         &self,
         connection: Connection,
@@ -140,7 +139,7 @@ impl<R: DeserializeOwned + Send + 'static> ProtocolHandler for IrohProtocol<R> {
         let request_id = self
             .request_id
             .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
-        let fut = handle_connection(connection, handler).map_err(AcceptError::from_err);
+        let fut = handle_connection::<S>(connection, handler).map_err(AcceptError::from_err);
         let span = trace_span!("rpc", id = request_id);
         Box::pin(fut.instrument(span))
     }
@@ -151,8 +150,8 @@ impl<R: DeserializeOwned + Send + 'static> ProtocolHandler for IrohProtocol<R> {
 /// Can be added to an [`iroh::protocol::Router`] to handle incoming connections for an ALPN string.
 ///
 /// For details about when it is safe to use 0rtt, see https://www.iroh.computer/blog/0rtt-api
-pub struct Iroh0RttProtocol<R> {
-    handler: Handler<R>,
+pub struct Iroh0RttProtocol<S> {
+    handler: Handler<S>,
     request_id: AtomicU64,
 }
 
@@ -162,25 +161,25 @@ impl<T> fmt::Debug for Iroh0RttProtocol<T> {
     }
 }
 
-impl<R: DeserializeOwned + Send + 'static> Iroh0RttProtocol<R> {
-    pub fn with_sender(local_sender: impl Into<LocalSender<R>>) -> Self
+impl<S: Service> Iroh0RttProtocol<S> {
+    pub fn with_sender(local_sender: impl Into<LocalSender<S>>) -> Self
     where
-        R: RemoteService,
+        S: RemoteService,
     {
-        let handler = R::remote_handler(local_sender.into());
+        let handler = S::remote_handler(local_sender.into());
         Self::new(handler)
     }
 
     /// Creates a new [`Iroh0RttProtocol`] for the `handler`.
-    pub fn new(handler: Handler<R>) -> Self {
+    pub fn new(handler: Handler<S>) -> Self {
         Self {
             handler,
             request_id: Default::default(),
         }
     }
 }
 
-impl<R: DeserializeOwned + Send + 'static> ProtocolHandler for Iroh0RttProtocol<R> {
+impl<S: Service> ProtocolHandler for Iroh0RttProtocol<S> {
     async fn on_connecting(&self, connecting: Connecting) -> Result<Connection, AcceptError> {
         let (conn, _zero_rtt_accepted) = connecting
             .into_0rtt()
@@ -196,29 +195,34 @@ impl<R: DeserializeOwned + Send + 'static> ProtocolHandler for Iroh0RttProtocol<
         let request_id = self
             .request_id
             .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
-        let fut = handle_connection(connection, handler).map_err(AcceptError::from_err);
+        let fut = handle_connection::<S>(connection, handler).map_err(AcceptError::from_err);
         let span = trace_span!("rpc", id = request_id);
         Box::pin(fut.instrument(span))
     }
 }
 
 /// Handles a single iroh connection with the provided `handler`.
-pub async fn handle_connection<R: DeserializeOwned + 'static>(
+///
+/// The wire format used depends on `S::SPAN_PROPAGATION` - if true, span context is expected.
+pub async fn handle_connection<S: Service>(
     connection: Connection,
-    handler: Handler<R>,
+    handler: Handler<S>,
 ) -> io::Result<()> {
     if let Ok(remote) = connection.remote_id() {
         tracing::Span::current().record("remote", tracing::field::display(remote.fmt_short()));
     }
     debug!("connection accepted");
     loop {
-        let Some((msg, rx, tx)) = read_request_raw(&connection).await? else {
+        let Some((msg, rx, tx)) = read_request_raw::<S>(&connection).await? else {
             return Ok(());
         };
         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`.
 pub async fn read_request<S: RemoteService>(
     connection: &Connection,
 ) -> std::io::Result<Option<S::Message>> {
@@ -231,12 +235,16 @@ pub async fn read_request<S: RemoteService>(
 ///
 /// This accepts a bi-directional stream from the connection and reads and parses the request.
 ///
+/// The wire format used depends on `S::SPAN_PROPAGATION`:
+/// - When `true`: expects `(Option<SpanContextCarrier>, Message)` tuple format
+/// - When `false`: expects plain `Message` format
+///
 /// Returns the parsed request and the stream pair if reading and parsing the request succeeded.
 /// Returns None if the remote closed the connection with error code `0`.
 /// Returns an error for all other failure cases.
-pub async fn read_request_raw<R: DeserializeOwned + 'static>(
+pub async fn read_request_raw<S: Service>(
     connection: &Connection,
-) -> std::io::Result<Option<(R, RecvStream, SendStream)>> {
+) -> std::io::Result<Option<(S, RecvStream, SendStream)>> {
     let (send, mut recv) = match connection.accept_bi().await {
         Ok((s, r)) => (s, r),
         Err(ConnectionError::ApplicationClosed(cause)) if cause.error_code.into_inner() == 0 => {
@@ -264,24 +272,31 @@ pub async fn read_request_raw<R: DeserializeOwned + 'static>(
         .await
         .map_err(|e| io::Error::new(io::ErrorKind::UnexpectedEof, e))?;
 
-    // Deserialize the payload which includes optional span context
-    // irpc-iroh uses irpc with default features, which include spans and rpc,
-    // so span_propagation module always exists
-    let (span_ctx, msg): (Option<irpc::span_propagation::SpanContextCarrier>, R) =
-        postcard::from_bytes(&buf).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+    // Deserialize based on S::SPAN_PROPAGATION
+    let msg: S = if S::SPAN_PROPAGATION {
+        let (span_ctx, msg): (Option<irpc::span_propagation::SpanContextCarrier>, S) =
+            postcard::from_bytes(&buf)
+                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
 
-    // Store span context in thread-local for use by with_remote_channels
-    if let Some(ctx) = span_ctx {
-        ctx.store_in_thread_local();
-    }
+        // Store span context in thread-local for use by with_remote_channels
+        if let Some(ctx) = span_ctx {
+            ctx.store_in_thread_local();
+        }
+
+        msg
+    } else {
+        postcard::from_bytes(&buf).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
+    };
 
     let rx = recv;
     let tx = send;
     Ok(Some((msg, rx, tx)))
 }
 
-/// Utility function to listen for incoming connections and handle them with the provided handler
-pub async fn listen<R: DeserializeOwned + 'static>(endpoint: iroh::Endpoint, handler: Handler<R>) {
+/// 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<S: Service>(endpoint: iroh::Endpoint, handler: Handler<S>) {
     let mut request_id = 0u64;
     let mut tasks = n0_future::task::JoinSet::new();
     loop {
@@ -300,7 +315,7 @@ pub async fn listen<R: DeserializeOwned + 'static>(endpoint: iroh::Endpoint, han
         let handler = handler.clone();
         let fut = async move {
             match incoming.await {
-                Ok(connection) => match handle_connection(connection, handler).await {
+                Ok(connection) => match handle_connection::<S>(connection, handler).await {
                     Err(err) => warn!("connection closed with error: {err:?}"),
                     Ok(()) => debug!("connection closed"),
                 },
diff --git a/src/lib.rs b/src/lib.rs
