diff --git a/tarpc/src/client.rs b/tarpc/src/client.rs index 9b543349..46bf562f 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) -> 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 cb69f680..1776a74a 100644 --- a/tarpc/src/client/in_flight_requests.rs +++ b/tarpc/src/client/in_flight_requests.rs @@ -77,20 +77,18 @@ impl InFlightRequests { } /// 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 { 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.