diff --git a/ddtrace/contrib/internal/grpc/aio_client_interceptor.py b/ddtrace/contrib/internal/grpc/aio_client_interceptor.py index 52d3aba9f61..b9852e27a10 100644 --- a/ddtrace/contrib/internal/grpc/aio_client_interceptor.py +++ b/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) diff --git a/releasenotes/notes/fix-grpc-aio-span-lifecycle-e9fd5a013b6f3d6f.yaml b/releasenotes/notes/fix-grpc-aio-span-lifecycle-e9fd5a013b6f3d6f.yaml new file mode 100644 index 00000000000..c576aad2d6e --- /dev/null +++ b/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. diff --git a/tests/contrib/grpc_aio/test_grpc_aio.py b/tests/contrib/grpc_aio/test_grpc_aio.py index fe473ea692e..6f041e7f26a 100644 --- a/tests/contrib/grpc_aio/test_grpc_aio.py +++ b/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 )