fix(grpc): finish aio client spans on abandoned calls

ddtrace/contrib/internal/grpc/aio_client_interceptor.py

@@ -224,7 +224,6 @@ async def _wrap_stream_response(
         span: Span,
     ) -> ResponseIterableType:
         try:
-            _handle_add_callback(call, _done_callback_stream(span))
             async for response in call:
                 yield response
         except StopAsyncIteration:
@@ -264,6 +263,9 @@ async def _wrap_unary_response(
             # So we can't handle the error in done callbacks.
             _handle_rpc_error(span, rpc_error)
             raise
+        except asyncio.CancelledError:
+            span.finish()
+            raise
 
 
 class _UnaryUnaryClientInterceptor(aio.UnaryUnaryClientInterceptor, _ClientInterceptor):
@@ -293,6 +295,7 @@ async def intercept_unary_stream(
             client_call_details,
         )
         call = await continuation(client_call_details, request)
+        _handle_add_callback(call, _done_callback_stream(span))
         return self._wrap_stream_response(call, span)
 
 
@@ -323,4 +326,5 @@ async def intercept_stream_stream(
             client_call_details,
         )
         call = await continuation(client_call_details, request_iterator)
+        _handle_add_callback(call, _done_callback_stream(span))
         return self._wrap_stream_response(call, span)

releasenotes/notes/fix-grpc-aio-span-lifecycle-e9fd5a013b6f3d6f.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    grpc: Fixes an issue where cancelled or unconsumed asynchronous client calls retain unfinished tracing spans.

tests/contrib/grpc_aio/test_grpc_aio.py

@@ -10,6 +10,8 @@
 from ddtrace.constants import ERROR_MSG
 from ddtrace.constants import ERROR_STACK
 from ddtrace.constants import ERROR_TYPE
+from ddtrace.contrib.internal.grpc.aio_client_interceptor import _StreamStreamClientInterceptor
+from ddtrace.contrib.internal.grpc.aio_client_interceptor import _UnaryStreamClientInterceptor
 from ddtrace.contrib.internal.grpc.patch import patch
 from ddtrace.contrib.internal.grpc.patch import unpatch
 from ddtrace.contrib.internal.grpc.utils import _parse_rpc_repr_string
@@ -38,6 +40,9 @@ async def SayHello(self, request, context):
             message = ";".join(w.key + "=" + w.value for w in metadata if w.key.startswith("x-datadog"))
             return HelloReply(message=message)
 
+        if request.name == "slow":
+            await asyncio.sleep(1)
+
         if request.name == "exception":
             await context.abort(grpc.StatusCode.INVALID_ARGUMENT, "abort_details")
 
@@ -164,6 +169,17 @@ def add_done_callback(self, unused_callback):
         pass
 
 
+class _CompletedStreamCall:
+    def add_done_callback(self, callback):
+        callback(self)
+
+    def done(self):
+        return True
+
+    def __repr__(self):
+        return 'status = StatusCode.OK, details = "complete"'
+
+
 @pytest.fixture(autouse=True)
 def patch_grpc_aio():
     patch()
@@ -401,6 +417,18 @@ async def test_unary_cancellation(server_info, tracer):
     assert len(spans) == 0
 
 
+@pytest.mark.parametrize("server_info", [_CoroHelloServicer()], indirect=True)
+async def test_unary_timeout_finishes_client_span(server_info, tracer):
+    async with aio.insecure_channel(server_info.target) as channel:
+        stub = HelloStub(channel)
+        with pytest.raises(asyncio.TimeoutError):
+            await asyncio.wait_for(stub.SayHello(HelloRequest(name="slow")), timeout=0.01)
+
+    await asyncio.sleep(0.1)
+    client_spans = [span for span in _get_spans(tracer) if span.service == "grpc-aio-client"]
+    assert len(client_spans) == 1
+
+
 @pytest.mark.parametrize(
     "server_info", [_CoroHelloServicer(), _AsyncGenHelloServicer(), _SyncHelloServicer()], indirect=True
 )
@@ -430,6 +458,37 @@ async def test_server_streaming(server_info, tracer):
     _check_server_span(server_span, "grpc-aio-server", "SayHelloTwice", "server_streaming")
 
 
+@pytest.mark.parametrize(
+    "interceptor,intercept_method,request_arg",
+    [
+        (_UnaryStreamClientInterceptor("localhost", 50051), "intercept_unary_stream", HelloRequest(name="test")),
+        (
+            _StreamStreamClientInterceptor("localhost", 50051),
+            "intercept_stream_stream",
+            iter([HelloRequest(name="test")]),
+        ),
+    ],
+    ids=["server_streaming", "bidi_streaming"],
+)
+async def test_streaming_uniterated_finishes_client_span(tracer, interceptor, intercept_method, request_arg):
+    call = _CompletedStreamCall()
+
+    async def continuation(client_call_details, request):
+        return call
+
+    client_call_details = aio.ClientCallDetails(
+        b"/helloworld.Hello/SayHelloTwice",
+        None,
+        None,
+        None,
+        None,
+    )
+    await getattr(interceptor, intercept_method)(continuation, client_call_details, request_arg)
+
+    client_spans = [span for span in _get_spans(tracer) if span.service == "grpc-aio-client"]
+    assert len(client_spans) == 1
+
+
 @pytest.mark.parametrize(
     "server_info", [_CoroHelloServicer(), _AsyncGenHelloServicer(), _SyncHelloServicer()], indirect=True
 )