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

@@ -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
     }
 }
 

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.