Repository: open-telemetry/opentelemetry-python-contrib

Issue:
gRPC server instrumentation creates multiple traces on streaming requests
**Environment**
Current `master` code, basically the sample code in the documentation, testing with a unary request vs. a streaming request.

**Steps to reproduce**
Create a simple gRPC servicer with two RPCs, one which returns a single message (the unary response), and one which yields items in a list for a streaming response.

The key here is to make an instrumented request within the primary request handler (I'm using a simple HTTP get with the Requests instrumentation), so you get an _additional_ span which should be attached to the same trace.

**What is the expected behavior?**
A single trace with the main span, and a second child span for the HTTP request.

**What is the actual behavior?**
Two separate traces, each containing a single span.

**Additional context**
The problem _only_ occurs on streaming requests - I'm sure the reworking I did as part of https://github.com/open-telemetry/opentelemetry-python/pull/1171 is where the problem started, I didn't take into account the streaming case specifically with multiple spans, and naturally, there are no tests for anything streaming, only unary responses.

So as part of this, we'll need some useful tests as well.  I'll see if I can write up my test case as an actual test case.

And again, I've got a vested interest in this working, so I'll have a PR up soon.


Discussion:

Resolving patch:
diff --git a/instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_server.py b/instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_server.py
--- a/instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_server.py
+++ b/instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_server.py
@@ -239,6 +239,15 @@ def intercept_service(self, continuation, handler_call_details):
         def telemetry_wrapper(behavior, request_streaming, response_streaming):
             def telemetry_interceptor(request_or_iterator, context):
 
+                # handle streaming responses specially
+                if response_streaming:
+                    return self._intercept_server_stream(
+                        behavior,
+                        handler_call_details,
+                        request_or_iterator,
+                        context,
+                    )
+
                 with self._set_remote_context(context):
                     with self._start_span(
                         handler_call_details, context
@@ -249,6 +258,7 @@ def telemetry_interceptor(request_or_iterator, context):
                         # And now we run the actual RPC.
                         try:
                             return behavior(request_or_iterator, context)
+
                         except Exception as error:
                             # Bare exceptions are likely to be gRPC aborts, which
                             # we handle in our context wrapper.
@@ -263,3 +273,23 @@ def telemetry_interceptor(request_or_iterator, context):
         return _wrap_rpc_behavior(
             continuation(handler_call_details), telemetry_wrapper
         )
+
+    # Handle streaming responses separately - we have to do this
+    # to return a *new* generator or various upstream things
+    # get confused, or we'll lose the consistent trace
+    def _intercept_server_stream(
+        self, behavior, handler_call_details, request_or_iterator, context
+    ):
+
+        with self._set_remote_context(context):
+            with self._start_span(handler_call_details, context) as span:
+                context = _OpenTelemetryServicerContext(context, span)
+
+                try:
+                    yield from behavior(request_or_iterator, context)
+
+                except Exception as error:
+                    # pylint:disable=unidiomatic-typecheck
+                    if type(error) != Exception:
+                        span.record_exception(error)
+                    raise error


Test patch:
diff --git a/instrumentation/opentelemetry-instrumentation-grpc/tests/test_server_interceptor.py b/instrumentation/opentelemetry-instrumentation-grpc/tests/test_server_interceptor.py
--- a/instrumentation/opentelemetry-instrumentation-grpc/tests/test_server_interceptor.py
+++ b/instrumentation/opentelemetry-instrumentation-grpc/tests/test_server_interceptor.py
@@ -30,6 +30,12 @@
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.trace.status import StatusCode
 
+from .protobuf.test_server_pb2 import Request, Response
+from .protobuf.test_server_pb2_grpc import (
+    GRPCTestServerServicer,
+    add_GRPCTestServerServicer_to_server,
+)
+
 
 class UnaryUnaryMethodHandler(grpc.RpcMethodHandler):
     def __init__(self, handler):
@@ -51,6 +57,23 @@ def service(self, handler_call_details):
         return UnaryUnaryMethodHandler(self._unary_unary_handler)
 
 
+class Servicer(GRPCTestServerServicer):
+    """Our test servicer"""
+
+    # pylint:disable=C0103
+    def SimpleMethod(self, request, context):
+        return Response(
+            server_id=request.client_id, response_data=request.request_data,
+        )
+
+    # pylint:disable=C0103
+    def ServerStreamingMethod(self, request, context):
+        for data in ("one", "two", "three"):
+            yield Response(
+                server_id=request.client_id, response_data=data,
+            )
+
+
 class TestOpenTelemetryServerInterceptor(TestBase):
     def test_instrumentor(self):
         def handler(request, context):
@@ -134,25 +157,146 @@ def test_create_span(self):
         # Intercept gRPC calls...
         interceptor = server_interceptor()
 
-        # No-op RPC handler
-        def handler(request, context):
-            return b""
+        server = grpc.server(
+            futures.ThreadPoolExecutor(max_workers=1),
+            options=(("grpc.so_reuseport", 0),),
+            interceptors=[interceptor],
+        )
+        add_GRPCTestServerServicer_to_server(Servicer(), server)
+        port = server.add_insecure_port("[::]:0")
+        channel = grpc.insecure_channel("localhost:{:d}".format(port))
+
+        rpc_call = "/GRPCTestServer/SimpleMethod"
+        request = Request(client_id=1, request_data="test")
+        msg = request.SerializeToString()
+        try:
+            server.start()
+            channel.unary_unary(rpc_call)(msg)
+        finally:
+            server.stop(None)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+        span = spans_list[0]
+
+        self.assertEqual(span.name, rpc_call)
+        self.assertIs(span.kind, trace.SpanKind.SERVER)
+
+        # Check version and name in span's instrumentation info
+        self.check_span_instrumentation_info(
+            span, opentelemetry.instrumentation.grpc
+        )
 
+        # Check attributes
+        self.assert_span_has_attributes(
+            span,
+            {
+                "net.peer.ip": "[::1]",
+                "net.peer.name": "localhost",
+                "rpc.method": "SimpleMethod",
+                "rpc.service": "GRPCTestServer",
+                "rpc.system": "grpc",
+                "rpc.grpc.status_code": grpc.StatusCode.OK.value[0],
+            },
+        )
+
+    def test_create_two_spans(self):
+        """Verify that the interceptor captures sub spans within the given
+        trace"""
+
+        class TwoSpanServicer(GRPCTestServerServicer):
+            # pylint:disable=C0103
+            def SimpleMethod(self, request, context):
+
+                # create another span
+                tracer = trace.get_tracer(__name__)
+                with tracer.start_as_current_span("child") as child:
+                    child.add_event("child event")
+
+                return Response(
+                    server_id=request.client_id,
+                    response_data=request.request_data,
+                )
+
+        # Intercept gRPC calls...
+        interceptor = server_interceptor()
+
+        # setup the server
         server = grpc.server(
             futures.ThreadPoolExecutor(max_workers=1),
             options=(("grpc.so_reuseport", 0),),
             interceptors=[interceptor],
         )
+        add_GRPCTestServerServicer_to_server(TwoSpanServicer(), server)
+        port = server.add_insecure_port("[::]:0")
+        channel = grpc.insecure_channel("localhost:{:d}".format(port))
 
-        server.add_generic_rpc_handlers((UnaryUnaryRpcHandler(handler),))
+        # setup the RPC
+        rpc_call = "/GRPCTestServer/SimpleMethod"
+        request = Request(client_id=1, request_data="test")
+        msg = request.SerializeToString()
+        try:
+            server.start()
+            channel.unary_unary(rpc_call)(msg)
+        finally:
+            server.stop(None)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 2)
+        child_span = spans_list[0]
+        parent_span = spans_list[1]
+
+        self.assertEqual(parent_span.name, rpc_call)
+        self.assertIs(parent_span.kind, trace.SpanKind.SERVER)
+
+        # Check version and name in span's instrumentation info
+        self.check_span_instrumentation_info(
+            parent_span, opentelemetry.instrumentation.grpc
+        )
 
+        # Check attributes
+        self.assert_span_has_attributes(
+            parent_span,
+            {
+                "net.peer.ip": "[::1]",
+                "net.peer.name": "localhost",
+                "rpc.method": "SimpleMethod",
+                "rpc.service": "GRPCTestServer",
+                "rpc.system": "grpc",
+                "rpc.grpc.status_code": grpc.StatusCode.OK.value[0],
+            },
+        )
+
+        # Check the child span
+        self.assertEqual(child_span.name, "child")
+        self.assertEqual(
+            parent_span.context.trace_id, child_span.context.trace_id
+        )
+
+    def test_create_span_streaming(self):
+        """Check that the interceptor wraps calls with spans server-side, on a
+        streaming call."""
+
+        # Intercept gRPC calls...
+        interceptor = server_interceptor()
+
+        # setup the server
+        server = grpc.server(
+            futures.ThreadPoolExecutor(max_workers=1),
+            options=(("grpc.so_reuseport", 0),),
+            interceptors=[interceptor],
+        )
+        add_GRPCTestServerServicer_to_server(Servicer(), server)
         port = server.add_insecure_port("[::]:0")
         channel = grpc.insecure_channel("localhost:{:d}".format(port))
 
-        rpc_call = "TestServicer/handler"
+        # setup the RPC
+        rpc_call = "/GRPCTestServer/ServerStreamingMethod"
+        request = Request(client_id=1, request_data="test")
+        msg = request.SerializeToString()
         try:
             server.start()
-            channel.unary_unary(rpc_call)(b"")
+            list(channel.unary_stream(rpc_call)(msg))
         finally:
             server.stop(None)
 
@@ -174,13 +318,86 @@ def handler(request, context):
             {
                 "net.peer.ip": "[::1]",
                 "net.peer.name": "localhost",
-                "rpc.method": "handler",
-                "rpc.service": "TestServicer",
+                "rpc.method": "ServerStreamingMethod",
+                "rpc.service": "GRPCTestServer",
+                "rpc.system": "grpc",
+                "rpc.grpc.status_code": grpc.StatusCode.OK.value[0],
+            },
+        )
+
+    def test_create_two_spans_streaming(self):
+        """Verify that the interceptor captures sub spans in a
+        streaming call, within the given trace"""
+
+        class TwoSpanServicer(GRPCTestServerServicer):
+            # pylint:disable=C0103
+            def ServerStreamingMethod(self, request, context):
+
+                # create another span
+                tracer = trace.get_tracer(__name__)
+                with tracer.start_as_current_span("child") as child:
+                    child.add_event("child event")
+
+                for data in ("one", "two", "three"):
+                    yield Response(
+                        server_id=request.client_id, response_data=data,
+                    )
+
+        # Intercept gRPC calls...
+        interceptor = server_interceptor()
+
+        # setup the server
+        server = grpc.server(
+            futures.ThreadPoolExecutor(max_workers=1),
+            options=(("grpc.so_reuseport", 0),),
+            interceptors=[interceptor],
+        )
+        add_GRPCTestServerServicer_to_server(TwoSpanServicer(), server)
+        port = server.add_insecure_port("[::]:0")
+        channel = grpc.insecure_channel("localhost:{:d}".format(port))
+
+        # setup the RPC
+        rpc_call = "/GRPCTestServer/ServerStreamingMethod"
+        request = Request(client_id=1, request_data="test")
+        msg = request.SerializeToString()
+        try:
+            server.start()
+            list(channel.unary_stream(rpc_call)(msg))
+        finally:
+            server.stop(None)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 2)
+        child_span = spans_list[0]
+        parent_span = spans_list[1]
+
+        self.assertEqual(parent_span.name, rpc_call)
+        self.assertIs(parent_span.kind, trace.SpanKind.SERVER)
+
+        # Check version and name in span's instrumentation info
+        self.check_span_instrumentation_info(
+            parent_span, opentelemetry.instrumentation.grpc
+        )
+
+        # Check attributes
+        self.assert_span_has_attributes(
+            parent_span,
+            {
+                "net.peer.ip": "[::1]",
+                "net.peer.name": "localhost",
+                "rpc.method": "ServerStreamingMethod",
+                "rpc.service": "GRPCTestServer",
                 "rpc.system": "grpc",
                 "rpc.grpc.status_code": grpc.StatusCode.OK.value[0],
             },
         )
 
+        # Check the child span
+        self.assertEqual(child_span.name, "child")
+        self.assertEqual(
+            parent_span.context.trace_id, child_span.context.trace_id
+        )
+
     def test_span_lifetime(self):
         """Check that the span is active for the duration of the call."""