Bugfix for open-telemetry#257 - properly stream responses

Fixes open-telemetry#257

instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/__init__.py

@@ -144,6 +144,9 @@ def serve():
 # pylint:disable=unused-argument
 # isort:skip
 
+import logging
+log = logging.getLogger(__name__)
+
 
 class GrpcInstrumentorServer(BaseInstrumentor):
     """
@@ -184,30 +187,44 @@ class GrpcInstrumentorClient(BaseInstrumentor):
         grpc_client_instrumentor = GrpcInstrumentorClient()
         grpc.client_instrumentor.instrument()
 
+    Instrumetor arguments:
+        wrap_secure (bool): False to disable wrapping secure channels
+        wrap_insecure (bool): False to disable wrapping insecure channels
+        exporter: OpenTelemetry metrics exporter
+        interval (int): metrics export interval
+
     """
 
     def _instrument(self, **kwargs):
         exporter = kwargs.get("exporter", None)
         interval = kwargs.get("interval", 30)
-        if kwargs.get("channel_type") == "secure":
+
+        # preserve the old argument
+        if "wrap_secure" not in kwargs and kwargs.get("channel_type", "") == "secure":
+            kwargs["wrap_secure"] = True
+            kwargs["wrap_insecure"] = False
+
+        if kwargs.get("wrap_secure", True):
+            log.info("wrapping secure channels")
             _wrap(
                 "grpc",
                 "secure_channel",
                 partial(self.wrapper_fn, exporter, interval),
             )
-
-        else:
+        if kwargs.get("wrap_insecure", True):
+            log.info("wrapping insecure channels")
             _wrap(
                 "grpc",
                 "insecure_channel",
                 partial(self.wrapper_fn, exporter, interval),
             )
 
     def _uninstrument(self, **kwargs):
-        if kwargs.get("channel_type") == "secure":
+        if kwargs.get("wrap_secure", True):
             unwrap(grpc, "secure_channel")
 
-        else:
+        #else:
+        if kwargs.get("wrap_insecure", True):
             unwrap(grpc, "insecure_channel")
 
     def wrapper_fn(

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