allow server to mutate shared context

axos88 · axos88 · commit d3fe5c042b06 · 2025-11-25T16:05:33.000+01:00
diff --git a/tarpc/src/client.rs b/tarpc/src/client.rs
@@ -166,14 +166,19 @@ where
             })
             .await
             .map_err(|mpsc::error::SendError(_)| RpcError::Shutdown)?;
-        response_guard.response().await
+
+        let (response_ctx, r) = response_guard.response().await?;
+
+        ctx.shared_context = response_ctx.shared_context;
+
+        Ok(r)
     }
 }
 
 /// A server response that is completed by request dispatch when the corresponding response
 /// arrives off the wire.
 struct ResponseGuard<'a, Resp> {
-    response: &'a mut oneshot::Receiver<Result<Resp, RpcError>>,
+    response: &'a mut oneshot::Receiver<Result<(ClientContext, Resp), RpcError>>,
     cancellation: &'a RequestCancellation,
     request_id: u64,
     cancel: bool,
@@ -201,7 +206,7 @@ pub enum RpcError {
 }
 
 impl<Resp> ResponseGuard<'_, Resp> {
-    async fn response(mut self) -> Result<Resp, RpcError> {
+    async fn response(mut self) -> Result<(ClientContext, Resp), RpcError> {
         let response = (&mut self.response).await;
         // Cancel drop logic once a response has been received.
         self.cancel = false;
@@ -280,7 +285,7 @@ pub struct RequestDispatch<Req, Resp, C> {
     /// Requests that were dropped.
     canceled_requests: CanceledRequests,
     /// Requests already written to the wire that haven't yet received responses.
-    in_flight_requests: InFlightRequests<Result<Resp, RpcError>>,
+    in_flight_requests: InFlightRequests<Resp>,
     /// Configures limits to prevent unlimited resource usage.
     config: Config,
     /// Produces errors that can be sent in response to any unprocessed requests at the time
@@ -296,7 +301,7 @@ where
 {
     fn in_flight_requests<'a>(
         self: &'a mut Pin<&mut Self>,
-    ) -> &'a mut InFlightRequests<Result<Resp, RpcError>> {
+    ) -> &'a mut InFlightRequests<Resp> {
         self.as_mut().project().in_flight_requests
     }
 
@@ -522,12 +527,10 @@ where
         let trace_context = ctx.trace_context;
         let deadline = ctx.deadline;
 
-        let client_context = context::ClientContext::new(ctx);
-
         let request = ClientMessage::Request(Request {
             id: request_id,
             message: request,
-            context: client_context,
+            context: ClientContext::new(ctx),
         });
 
         self.in_flight_requests()
@@ -580,7 +583,7 @@ where
     fn complete(mut self: Pin<&mut Self>, response: Response<ClientContext, Resp>) -> bool {
         if let Some(span) = self.in_flight_requests().complete_request(
             response.request_id,
-            response.message.map_err(RpcError::Server),
+            response.message.map_err(RpcError::Server).map(|m| (response.context, m)),
         ) {
             let _entered = span.enter();
             tracing::debug!("ReceiveResponse");
@@ -688,11 +691,11 @@ where
 /// the lifecycle of the request.
 #[derive(Debug)]
 struct DispatchRequest<Req, Resp> {
-    pub ctx: context::SharedContext,
+    pub ctx: context::SharedContext, ///TODO: <-- this should be a &mut ClientContext
     pub span: Span,
     pub request_id: u64,
     pub request: Req,
-    pub response_completion: oneshot::Sender<Result<Resp, RpcError>>,
+    pub response_completion: oneshot::Sender<Result<(ClientContext, Resp), RpcError>>,
 }
 
 #[cfg(test)]
@@ -752,7 +755,7 @@ mod tests {
             .await
             .unwrap();
         assert_matches!(dispatch.as_mut().poll(cx), Poll::Pending);
-        assert_matches!(rx.try_recv(), Ok(Ok(resp)) if resp == "Resp");
+        assert_matches!(rx.try_recv(), Ok(Ok((_, resp))) if resp == "Resp");
     }
 
     #[tokio::test]
@@ -774,12 +777,7 @@ mod tests {
     async fn dispatch_response_doesnt_cancel_after_complete() {
         let (cancellation, mut canceled_requests) = cancellations();
         let (tx, mut response) = oneshot::channel();
-        tx.send(Ok(Response {
-            request_id: 0,
-            context: ClientContext::current(),
-            message: Ok("well done"),
-        }))
-        .unwrap();
+        tx.send(Ok((ClientContext::current(), "well done"))).unwrap();
         // resp's drop() is run, but should not send a cancel message.
         ResponseGuard {
             response: &mut response,
@@ -1116,37 +1114,11 @@ mod tests {
         (Box::pin(dispatch), channel, server_channel)
     }
 
-    async fn reserve_for_send<'a>(
-        channel: &'a mut Channel<String, String>,
-        response_completion: oneshot::Sender<Result<String, RpcError>>,
-        response: &'a mut oneshot::Receiver<Result<String, RpcError>>,
-    ) -> impl FnOnce(&str) -> ResponseGuard<'a, String> {
-        let permit = channel.to_dispatch.reserve().await.unwrap();
-        |request| {
-            let request_id =
-                u64::try_from(channel.next_request_id.fetch_add(1, Ordering::Relaxed)).unwrap();
-            let request = DispatchRequest {
-                ctx: SharedContext::current(),
-                span: Span::current(),
-                request_id,
-                request: request.to_string(),
-                response_completion,
-            };
-            permit.send(request);
-            ResponseGuard {
-                response,
-                cancellation: &channel.cancellation,
-                request_id,
-                cancel: true,
-            }
-        }
-    }
-
     async fn send_request<'a>(
         channel: &'a mut Channel<String, String>,
         request: &str,
-        response_completion: oneshot::Sender<Result<String, RpcError>>,
-        response: &'a mut oneshot::Receiver<Result<String, RpcError>>,
+        response_completion: oneshot::Sender<Result<(ClientContext, String), RpcError>>,
+        response: &'a mut oneshot::Receiver<Result<(ClientContext, String), RpcError>>,
     ) -> ResponseGuard<'a, String> {
         let request_id =
             u64::try_from(channel.next_request_id.fetch_add(1, Ordering::Relaxed)).unwrap();
@@ -1167,6 +1139,32 @@ mod tests {
         response_guard
     }
 
+    async fn reserve_for_send<'a>(
+        channel: &'a mut Channel<String, String>,
+        response_completion: oneshot::Sender<Result<(ClientContext, String), RpcError>>,
+        response: &'a mut oneshot::Receiver<Result<(ClientContext, String), RpcError>>,
+    ) -> impl FnOnce(&str) -> ResponseGuard<'a, String> {
+        let permit = channel.to_dispatch.reserve().await.unwrap();
+        |request| {
+            let request_id =
+              u64::try_from(channel.next_request_id.fetch_add(1, Ordering::Relaxed)).unwrap();
+            let request = DispatchRequest {
+                ctx: SharedContext::current(),
+                span: Span::current(),
+                request_id,
+                request: request.to_string(),
+                response_completion,
+            };
+            permit.send(request);
+            ResponseGuard {
+                response,
+                cancellation: &channel.cancellation,
+                request_id,
+                cancel: true,
+            }
+        }
+    }
+
     async fn send_response(
         channel: &mut UnboundedChannel<
             ClientMessage<ClientContext, String>,
diff --git a/tarpc/src/client/in_flight_requests.rs b/tarpc/src/client/in_flight_requests.rs
@@ -1,7 +1,4 @@
-use crate::{
-    trace,
-    util::{Compact, TimeUntil},
-};
+use crate::{trace, util::{Compact, TimeUntil}};
 use fnv::FnvHashMap;
 use std::time::Instant;
 use std::{
@@ -11,6 +8,8 @@ use std::{
 use tokio::sync::oneshot;
 use tokio_util::time::delay_queue::{self, DelayQueue};
 use tracing::Span;
+use crate::client::RpcError;
+use crate::context::ClientContext;
 
 /// Requests already written to the wire that haven't yet received responses.
 #[derive(Debug)]
@@ -32,7 +31,7 @@ impl<Resp> Default for InFlightRequests<Resp> {
 struct RequestData<Res> {
     ctx: trace::Context,
     span: Span,
-    response_completion: oneshot::Sender<Res>,
+    response_completion: oneshot::Sender<Result<(ClientContext, Res), RpcError>>,
     /// The key to remove the timer for the request's deadline.
     deadline_key: delay_queue::Key,
 }
@@ -60,7 +59,7 @@ impl<Res> InFlightRequests<Res> {
         ctx: trace::Context,
         deadline: Instant,
         span: Span,
-        response_completion: oneshot::Sender<Res>,
+        response_completion: oneshot::Sender<Result<(ClientContext, Res), RpcError>>,
     ) -> Result<(), AlreadyExistsError> {
         match self.request_data.entry(request_id) {
             hash_map::Entry::Vacant(vacant) => {
@@ -78,8 +77,8 @@ impl<Res> InFlightRequests<Res> {
         }
     }
 
-    /// Removes a request without aborting. Returns true iff the request was found.
-    pub fn complete_request(&mut self, request_id: u64, result: Res) -> Option<Span> {
+    /// Removes a request without aborting. Returns true if the request was found.
+    pub fn complete_request(&mut self, request_id: u64, result: Result<(ClientContext, Res), RpcError>) -> Option<Span> {
         if let Some(request_data) = self.request_data.remove(&request_id) {
             self.request_data.compact(0.1);
             self.deadlines.remove(&request_data.deadline_key);
@@ -97,7 +96,7 @@ impl<Res> InFlightRequests<Res> {
     /// Returns Spans for all completes requests.
     pub fn complete_all_requests<'a>(
         &'a mut self,
-        mut result: impl FnMut() -> Res + 'a,
+        mut result: impl FnMut() -> Result<(ClientContext, Res), RpcError> + 'a,
     ) -> impl Iterator<Item = Span> + 'a {
         self.deadlines.clear();
         self.request_data.drain().map(move |(_, request_data)| {
@@ -123,7 +122,7 @@ impl<Res> InFlightRequests<Res> {
     pub fn poll_expired(
         &mut self,
         cx: &mut Context,
-        expired_error: impl Fn() -> Res,
+        expired_error: impl Fn() -> Result<(ClientContext, Res), RpcError>,
     ) -> Poll<Option<u64>> {
         self.deadlines.poll_expired(cx).map(|expired| {
             let request_id = expired?.into_inner();
