fix(grpc): leaking spans when using gRPC future interface [backport 3.16] (#14900)

Backport 75a8246 from #14875 to 3.16.

We've tracked down some issues related to broken span hierarchy down to
the way that the gRPC instrumentation deals with entering spans, which
causes the grpc span to leak outside it's call when using the .future()
interface. This PR changes the instrumentation to explicitly only enter
the span when calling the continuations instead of just entering
immediately.

## Description

<!-- Provide an overview of the change and motivation for the change -->

## Testing

I've added tests, but haven't run them for a while and working on
getting it running locally. Will update if I run into any issues

## Risks

<!-- Note any risks associated with this change, or "None" if no risks
-->

## Additional Notes

Co-authored-by: Sam Schlegel <[email protected]>

Signed-off-by: Sam Schlegel <[email protected]>
Co-authored-by: Brett Langdon <[email protected]>
Co-authored-by: Sam Schlegel <[email protected]>
Co-authored-by: Christophe Papazian <[email protected]>

Author: 4 people

ddtrace/contrib/internal/grpc/client_interceptor.py

@@ -1,4 +1,5 @@
 import collections
+from contextlib import contextmanager
 
 import grpc
 import wrapt
@@ -20,6 +21,8 @@
 from ddtrace.internal.schema import schematize_url_operation
 from ddtrace.internal.schema.span_attribute_schema import SpanDirection
 from ddtrace.propagation.http import HTTPPropagator
+from ddtrace.trace import Span
+from ddtrace.trace import Tracer
 
 
 log = get_logger(__name__)
@@ -120,10 +123,21 @@ def _handle_error(span, response_error, status_code):
             status_code = str(response_error.code())
 
 
+@contextmanager
+def _activated_span(tracer: Tracer, span: Span):
+    prev_span = tracer.context_provider.active()
+    tracer.context_provider.activate(span)
+    try:
+        yield
+    finally:
+        tracer.context_provider.activate(prev_span)
+
+
 class _WrappedResponseCallFuture(wrapt.ObjectProxy):
-    def __init__(self, wrapped, span):
+    def __init__(self, wrapped, span, tracer):
         super(_WrappedResponseCallFuture, self).__init__(wrapped)
         self._span = span
+        self._tracer = tracer
         # Registers callback on the _MultiThreadedRendezvous future to finish
         # span in case StopIteration is never raised but RPC is terminated
         _handle_response(self._span, self.__wrapped__)
@@ -140,7 +154,8 @@ def _next(self):
         # https://github.com/googleapis/python-api-core/blob/35e87e0aca52167029784379ca84e979098e1d6c/google/api_core/grpc_helpers.py#L84
         # https://github.com/GoogleCloudPlatform/grpc-gcp-python/blob/5a2cd9807bbaf1b85402a2a364775e5b65853df6/src/grpc_gcp/_channel.py#L102
         try:
-            return next(self.__wrapped__)
+            with _activated_span(self._tracer, self._span):
+                return next(self.__wrapped__)
         except StopIteration:
             # Callback will handle span finishing
             raise
@@ -178,13 +193,17 @@ def __init__(self, pin, host, port):
         self._port = port
 
     def _intercept_client_call(self, method_kind, client_call_details):
-        tracer = self._pin.tracer
+        tracer: Tracer = self._pin.tracer
 
-        span = tracer.trace(
+        # Instead of using .trace, create the span and activate it at points where we call the continuations
+        # This avoids the issue of spans being leaked when using the .future interface.
+        parent = tracer.context_provider.active()
+        span = tracer.start_span(
             schematize_url_operation("grpc", protocol="grpc", direction=SpanDirection.OUTBOUND),
             span_type=SpanTypes.GRPC,
             service=trace_utils.ext_service(self._pin, config.grpc),
             resource=client_call_details.method,
+            child_of=parent,
         )
 
         span.set_tag_str(COMPONENT, config.grpc.integration_name)
@@ -206,7 +225,8 @@ def _intercept_client_call(self, method_kind, client_call_details):
         # propagate distributed tracing headers if available
         headers = {}
         if config.grpc.distributed_tracing_enabled:
-            HTTPPropagator.inject(span.context, headers)
+            # NOTE: We need to pass the span to the HTTPPropagator since it isn't active at this point
+            HTTPPropagator.inject(span.context, headers, span)
 
         metadata = []
         if client_call_details.metadata is not None:
@@ -227,15 +247,16 @@ def intercept_unary_unary(self, continuation, client_call_details, request):
             constants.GRPC_METHOD_KIND_UNARY,
             client_call_details,
         )
-        try:
-            response = continuation(client_call_details, request)
-            _handle_response(span, response)
-        except grpc.RpcError as rpc_error:
-            # DEV: grpcio<1.18.0 grpc.RpcError is raised rather than returned as response
-            # https://github.com/grpc/grpc/commit/8199aff7a66460fbc4e9a82ade2e95ef076fd8f9
-            # handle as a response
-            _handle_response(span, rpc_error)
-            raise
+        with _activated_span(self._pin.tracer, span):
+            try:
+                response = continuation(client_call_details, request)
+                _handle_response(span, response)
+            except grpc.RpcError as rpc_error:
+                # DEV: grpcio<1.18.0 grpc.RpcError is raised rather than returned as response
+                # https://github.com/grpc/grpc/commit/8199aff7a66460fbc4e9a82ade2e95ef076fd8f9
+                # handle as a response
+                _handle_response(span, rpc_error)
+                raise
 
         return response
 
@@ -244,24 +265,26 @@ def intercept_unary_stream(self, continuation, client_call_details, request):
             constants.GRPC_METHOD_KIND_SERVER_STREAMING,
             client_call_details,
         )
-        response_iterator = continuation(client_call_details, request)
-        response_iterator = _WrappedResponseCallFuture(response_iterator, span)
+        with _activated_span(self._pin.tracer, span):
+            response_iterator = continuation(client_call_details, request)
+            response_iterator = _WrappedResponseCallFuture(response_iterator, span, self._pin.tracer)
         return response_iterator
 
     def intercept_stream_unary(self, continuation, client_call_details, request_iterator):
         span, client_call_details = self._intercept_client_call(
             constants.GRPC_METHOD_KIND_CLIENT_STREAMING,
             client_call_details,
         )
-        try:
-            response = continuation(client_call_details, request_iterator)
-            _handle_response(span, response)
-        except grpc.RpcError as rpc_error:
-            # DEV: grpcio<1.18.0 grpc.RpcError is raised rather than returned as response
-            # https://github.com/grpc/grpc/commit/8199aff7a66460fbc4e9a82ade2e95ef076fd8f9
-            # handle as a response
-            _handle_response(span, rpc_error)
-            raise
+        with _activated_span(self._pin.tracer, span):
+            try:
+                response = continuation(client_call_details, request_iterator)
+                _handle_response(span, response)
+            except grpc.RpcError as rpc_error:
+                # DEV: grpcio<1.18.0 grpc.RpcError is raised rather than returned as response
+                # https://github.com/grpc/grpc/commit/8199aff7a66460fbc4e9a82ade2e95ef076fd8f9
+                # handle as a response
+                _handle_response(span, rpc_error)
+                raise
 
         return response
 
@@ -270,6 +293,7 @@ def intercept_stream_stream(self, continuation, client_call_details, request_ite
             constants.GRPC_METHOD_KIND_BIDI_STREAMING,
             client_call_details,
         )
-        response_iterator = continuation(client_call_details, request_iterator)
-        response_iterator = _WrappedResponseCallFuture(response_iterator, span)
+        with _activated_span(self._pin.tracer, span):
+            response_iterator = continuation(client_call_details, request_iterator)
+            response_iterator = _WrappedResponseCallFuture(response_iterator, span, self._pin.tracer)
         return response_iterator

releasenotes/notes/fix-grpc-tracing-leaking-span-context-when-using-future-interface-8a9cf0279263eb20.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    grpc: This fix resolves an issue where the internal span was left active in the caller when using the future interface.

tests/contrib/grpc/test_grpc.py

@@ -612,6 +612,43 @@ def test_schematized_operation_name_v1(self):
         assert client_span.name == "grpc.client.request", "Expected 'grpc.client.request', got %s" % client_span.name
         assert server_span.name == "grpc.server.request", "Expected 'grpc.server.request', got %s" % server_span.name
 
+    def test_span_parent_is_maintained(self):
+        with self.tracer.trace("root") as _:
+            with grpc.insecure_channel("localhost:%d" % (_GRPC_PORT)) as channel:
+                stub = HelloStub(channel)
+                _ = stub.SayHello(HelloRequest(name="test"))
+
+        spans = self.get_spans_with_sync_and_assert(size=3)
+        root = spans[0]
+        client = spans[1]
+        server = spans[2]
+        assert root.span_id == client.parent_id
+        assert client.span_id == server.parent_id
+
+    def test_active_span_doesnt_leak_with_future(self):
+        with grpc.insecure_channel("localhost:%d" % (_GRPC_PORT)) as channel:
+            stub = HelloStub(channel)
+            future = stub.SayHello.future(HelloRequest(name="test"))
+            assert self.tracer.current_span() is None
+            # wait so that we don't cancel the request
+            future.result()
+
+        self.get_spans_with_sync_and_assert(size=2)
+
+    def test_span_active_in_custom_interceptor(self):
+        # add an interceptor that raises a custom exception and check error tags
+        # are added to spans
+        interceptor = _SpanActivationClientInterceptor(self.tracer)
+        with grpc.insecure_channel("localhost:%d" % (_GRPC_PORT)) as channel:
+            intercept_channel = grpc.intercept_channel(channel, interceptor)
+            stub = HelloStub(intercept_channel)
+            stub.SayHello(HelloRequest(name="test"))
+
+        spans = self.get_spans_with_sync_and_assert(size=2)
+        client_span = spans[0]
+
+        assert client_span.get_tag("custom_interceptor_worked")
+
 
 class _CustomException(Exception):
     pass
@@ -686,3 +723,21 @@ def handler(request, context):
         channel.unary_unary("/pkg.Servicer/Handler")(b"request")
     finally:
         server.stop(None)
+
+
+class _SpanActivationClientInterceptor(grpc.UnaryUnaryClientInterceptor):
+    def __init__(self, tracer) -> None:
+        super().__init__()
+        self.tracer = tracer
+
+    def _intercept_call(self, continuation, client_call_details, request_or_iterator):
+        # allow computation to complete
+        span = self.tracer.current_span()
+        assert span is not None
+
+        span.set_tag("custom_interceptor_worked")
+
+        return continuation(client_call_details, request_or_iterator)
+
+    def intercept_unary_unary(self, continuation, client_call_details, request):
+        return self._intercept_call(continuation, client_call_details, request)