Wrap strean response in proxy that manage span lifetime

Author: UOndro

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

@@ -21,10 +21,11 @@
 
 import logging
 from collections import OrderedDict
-from functools import partial
 from typing import Callable, MutableMapping
 
 import grpc
+import wrapt
+
 from opentelemetry import context, trace
 from opentelemetry.instrumentation.grpc import grpcext
 from opentelemetry.instrumentation.grpc._utilities import RpcInfo
@@ -72,11 +73,59 @@ def _safe_invoke(function: Callable, *args):
             "Error when invoking function '%s'", function_name, exc_info=ex
         )
 
+
+class OpenTelemetryStreamWrapper(wrapt.ObjectProxy):
+    def __init__(self, wrapped, span: trace.Span):
+        super().__init__(wrapped)
+        self._self_span = span
+
+    def _end_span_if_not_already_ended(self, status_code=None, status=None):
+        if self._self_span.end_time is None:
+            self._self_span.end()
+            if status_code is not None:
+                self._self_span.set_attribute(
+                    SpanAttributes.RPC_GRPC_STATUS_CODE, status_code
+                )
+            if status is not None:
+                self._self_span.set_status(status)
+
+    def __del__(self):
+        self._end_span_if_not_already_ended()
+        self.__wrapped__.__del__()
+
+    def __iter__(self):
+        return self
+
+    def cancel(self):
+        self._end_span_if_not_already_ended(
+            status_code=grpc.StatusCode.CANCELLED.value[0]
+        )
+        return self.__wrapped__.cancel()
+
+    def __next__(self):
+        return self._next()
+
+    def next(self):
+        return self._next()
+
+    def _next(self):
+        try:
+            return self.__wrapped__._next()
+        except StopIteration:
+            self._end_span_if_not_already_ended()
+            raise
+        except grpc.RpcError as err:
+            self._end_span_if_not_already_ended(
+                err.code().value[0], Status(StatusCode.ERROR)
+            )
+            raise err
+
+
 class OpenTelemetryClientInterceptor(
     grpcext.UnaryClientInterceptor, grpcext.StreamClientInterceptor
 ):
     def __init__(
-            self, tracer, filter_=None, request_hook=None, response_hook=None
+        self, tracer, filter_=None, request_hook=None, response_hook=None
     ):
         self._tracer = tracer
         self._filter = filter_
@@ -130,10 +179,10 @@ def _intercept(self, request, metadata, client_info, invoker):
         else:
             mutable_metadata = OrderedDict(metadata)
         with self._start_span(
-                client_info.full_method,
-                end_on_exit=False,
-                record_exception=False,
-                set_status_on_exception=False,
+            client_info.full_method,
+            end_on_exit=False,
+            record_exception=False,
+            set_status_on_exception=False,
         ) as span:
             result = None
             try:
@@ -187,16 +236,15 @@ def intercept_unary(self, request, metadata, client_info, invoker):
     # the span across the generated responses and detect any errors, we wrap
     # the result in a new generator that yields the response values.
     def _intercept_server_stream(
-            self, request_or_iterator, metadata, client_info, invoker
+        self, request_or_iterator, metadata, client_info, invoker
     ):
         if not metadata:
             mutable_metadata = OrderedDict()
         else:
             mutable_metadata = OrderedDict(metadata)
 
         with self._start_span(
-                client_info.full_method,
-                end_on_exit=False
+            client_info.full_method, end_on_exit=False
         ) as span:
             inject(mutable_metadata, setter=_carrier_setter)
             metadata = tuple(mutable_metadata.items())
@@ -211,27 +259,10 @@ def _intercept_server_stream(
 
             stream = invoker(request_or_iterator, metadata)
 
-            def done_callback(future, span_):
-                try:
-                    future.result()
-                except grpc.FutureCancelledError:
-                    span_.set_status(Status(StatusCode.OK))
-                    span_.set_attribute(
-                        SpanAttributes.RPC_GRPC_STATUS_CODE, grpc.StatusCode.CANCELLED.value[0]
-                    )
-                except grpc.RpcError as err:
-                    span_.set_status(Status(StatusCode.ERROR))
-                    span_.set_attribute(
-                        SpanAttributes.RPC_GRPC_STATUS_CODE, err.code().value[0]
-                    )
-                finally:
-                    span_.end()
-
-            stream.add_done_callback(partial(done_callback, span_=span))
-            return stream
+            return OpenTelemetryStreamWrapper(stream, span)
 
     def intercept_stream(
-            self, request_or_iterator, metadata, client_info, invoker
+        self, request_or_iterator, metadata, client_info, invoker
     ):
         if context.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
             return invoker(request_or_iterator, metadata)

instrumentation/opentelemetry-instrumentation-grpc/tests/_client.py

@@ -53,7 +53,7 @@ def server_streaming_method(stub, error=False, serialize=True):
     return response_iterator
 
 
-def bidirectional_streaming_method(stub, error=False):
+def bidirectional_streaming_method(stub, error=False, serialize=True):
     def request_messages():
         for _ in range(5):
             request = Request(
@@ -62,5 +62,6 @@ def request_messages():
             yield request
 
     response_iterator = stub.BidirectionalStreamingMethod(request_messages())
-
-    list(response_iterator)
+    if serialize:
+        list(response_iterator)
+    return response_iterator

instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # pylint:disable=cyclic-import
-from time import sleep
 
 import grpc
 from tests.protobuf import (  # pylint: disable=no-name-in-module
@@ -56,32 +55,32 @@ def __init__(self):
         pass
 
     def intercept_unary_unary(
-            self, continuation, client_call_details, request
+        self, continuation, client_call_details, request
     ):
         return self._intercept_call(continuation, client_call_details, request)
 
     def intercept_unary_stream(
-            self, continuation, client_call_details, request
+        self, continuation, client_call_details, request
     ):
         return self._intercept_call(continuation, client_call_details, request)
 
     def intercept_stream_unary(
-            self, continuation, client_call_details, request_iterator
+        self, continuation, client_call_details, request_iterator
     ):
         return self._intercept_call(
             continuation, client_call_details, request_iterator
         )
 
     def intercept_stream_stream(
-            self, continuation, client_call_details, request_iterator
+        self, continuation, client_call_details, request_iterator
     ):
         return self._intercept_call(
             continuation, client_call_details, request_iterator
         )
 
     @staticmethod
     def _intercept_call(
-            continuation, client_call_details, request_or_iterator
+        continuation, client_call_details, request_or_iterator
     ):
         return continuation(client_call_details, request_or_iterator)
 
@@ -94,9 +93,7 @@ def setUp(self):
         self.server.start()
         # use a user defined interceptor along with the opentelemetry client interceptor
         interceptors = [Interceptor()]
-        self.channel = grpc.insecure_channel("localhost:25565", options=[
-            # (grpc.experimental.ChannelOptions.SingleThreadedUnaryStream, 1)
-        ])
+        self.channel = grpc.insecure_channel("localhost:25565")
         self.channel = grpc.intercept_channel(self.channel, *interceptors)
         self._stub = test_server_pb2_grpc.GRPCTestServerStub(self.channel)
 
@@ -173,14 +170,11 @@ def test_unary_stream(self):
         )
 
     def test_unary_stream_can_be_cancel(self):
-        responses = server_streaming_method(self._stub)
+        responses = server_streaming_method(self._stub, serialize=False)
         for i, _ in enumerate(responses):
             if i == 1:
                 responses.cancel()
                 break
-        sleep(10)
-        self.server.stop(None)
-        self.channel.close()
         spans = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans), 1)
         span = spans[0]
@@ -205,6 +199,33 @@ def test_unary_stream_can_be_cancel(self):
             },
         )
 
+    def test_finished_stream_cancel_does_not_change_status_of_span(self):
+        responses = server_streaming_method(self._stub, serialize=True)
+        responses.cancel()
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+
+        self.assertEqual(span.name, "/GRPCTestServer/ServerStreamingMethod")
+        self.assertIs(span.kind, trace.SpanKind.CLIENT)
+
+        # Check version and name in span's instrumentation info
+        self.assertEqualSpanInstrumentationInfo(
+            span, opentelemetry.instrumentation.grpc
+        )
+
+        self.assertSpanHasAttributes(
+            span,
+            {
+                SpanAttributes.RPC_METHOD: "ServerStreamingMethod",
+                SpanAttributes.RPC_SERVICE: "GRPCTestServer",
+                SpanAttributes.RPC_SYSTEM: "grpc",
+                SpanAttributes.RPC_GRPC_STATUS_CODE: grpc.StatusCode.OK.value[
+                    0
+                ],
+            },
+        )
+
     def test_stream_unary(self):
         client_streaming_method(self._stub)
         spans = self.memory_exporter.get_finished_spans()
@@ -259,6 +280,38 @@ def test_stream_stream(self):
             },
         )
 
+    def test_stream_stream_can_be_cancel(self):
+        responses = bidirectional_streaming_method(self._stub, serialize=False)
+        for i, _ in enumerate(responses):
+            if i == 1:
+                responses.cancel()
+                break
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+
+        self.assertEqual(
+            span.name, "/GRPCTestServer/BidirectionalStreamingMethod"
+        )
+        self.assertIs(span.kind, trace.SpanKind.CLIENT)
+
+        # Check version and name in span's instrumentation info
+        self.assertEqualSpanInstrumentationInfo(
+            span, opentelemetry.instrumentation.grpc
+        )
+
+        self.assertSpanHasAttributes(
+            span,
+            {
+                SpanAttributes.RPC_METHOD: "BidirectionalStreamingMethod",
+                SpanAttributes.RPC_SERVICE: "GRPCTestServer",
+                SpanAttributes.RPC_SYSTEM: "grpc",
+                SpanAttributes.RPC_GRPC_STATUS_CODE: grpc.StatusCode.CANCELLED.value[
+                    0
+                ],
+            },
+        )
+
     def test_error_simple(self):
         with self.assertRaises(grpc.RpcError):
             simple_method(self._stub, error=True)
@@ -308,7 +361,7 @@ def test_error_stream_stream(self):
         )
 
     def test_client_interceptor_trace_context_propagation(
-            self,
+        self,
     ):  # pylint: disable=no-self-use
         """ensure that client interceptor correctly inject trace context into all outgoing requests."""
         previous_propagator = get_global_textmap()

instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor_filter.py

@@ -98,9 +98,7 @@ def setUp(self):
         self.server.start()
         # use a user defined interceptor along with the opentelemetry client interceptor
         interceptors = [Interceptor()]
-        self.channel = grpc.insecure_channel("localhost:25565",options=[
-            (grpc.experimental.ChannelOptions.SingleThreadedUnaryStream, 1)
-        ])
+        self.channel = grpc.insecure_channel("localhost:25565")
         self.channel = grpc.intercept_channel(self.channel, *interceptors)
         self._stub = test_server_pb2_grpc.GRPCTestServerStub(self.channel)