From 84932df9b48b9b64e12517a9944ba13d181e08c4 Mon Sep 17 00:00:00 2001
From: Tim Kuehn <tikue@google.com>
Date: Mon, 9 Jan 2023 15:06:10 -0800
Subject: [PATCH] Return span in InFlightRequests::complete_request.

Rather than returning a bool, return the Span associated with the
request. This gives RequestDispatch more flexibility to annotate the
request span.
---
 tarpc/src/client.rs                    | 9 +++++++--
 tarpc/src/client/in_flight_requests.rs | 8 +++-----
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/tarpc/src/client.rs b/tarpc/src/client.rs
index 9b54334..46bf562 100644
--- a/tarpc/src/client.rs
+++ b/tarpc/src/client.rs
@@ -543,10 +543,15 @@ where
 
     /// Sends a server response to the client task that initiated the associated request.
     fn complete(mut self: Pin<&mut Self>, response: Response<Resp>) -> bool {
-        self.in_flight_requests().complete_request(
+        if let Some(span) = self.in_flight_requests().complete_request(
             response.request_id,
             response.message.map_err(RpcError::Server),
-        )
+        ) {
+            let _entered = span.enter();
+            tracing::info!("ReceiveResponse");
+            return true;
+        }
+        false
     }
 }
 
diff --git a/tarpc/src/client/in_flight_requests.rs b/tarpc/src/client/in_flight_requests.rs
index cb69f68..1776a74 100644
--- a/tarpc/src/client/in_flight_requests.rs
+++ b/tarpc/src/client/in_flight_requests.rs
@@ -77,20 +77,18 @@ 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) -> bool {
+    pub fn complete_request(&mut self, request_id: u64, result: Res) -> Option<Span> {
         if let Some(request_data) = self.request_data.remove(&request_id) {
-            let _entered = request_data.span.enter();
-            tracing::info!("ReceiveResponse");
             self.request_data.compact(0.1);
             self.deadlines.remove(&request_data.deadline_key);
             let _ = request_data.response_completion.send(result);
-            return true;
+            return Some(request_data.span);
         }
 
         tracing::debug!("No in-flight request found for request_id = {request_id}.");
 
         // If the response completion was absent, then the request was already canceled.
-        false
+        None
     }
 
     /// Completes all requests using the provided function.