def intercept_stream(self, request_or_iterator, metadata, client_info, invoker): if client_info.is_server_stream: return self._intercept_server_stream(request_or_iterator, metadata, client_info, invoker) if not metadata: mutable_metadata = OrderedDict() else: mutable_metadata = OrderedDict(metadata) with self._start_guarded_span(client_info.full_method) as guarded_span: _inject_span_context(mutable_metadata) metadata = tuple(mutable_metadata.items()) rpc_info = RpcInfo( full_method=client_info.full_method, metadata=metadata, timeout=client_info.timeout, request=request_or_iterator, ) rpc_info.request = request_or_iterator try: result = invoker(request_or_iterator, metadata) except grpc.RpcError as err: guarded_span.generated_span.set_status(Status( StatusCode.ERROR)) guarded_span.generated_span.set_attribute( "rpc.grpc.status_code", err.code().value[0], ) raise err return self._trace_result(guarded_span, rpc_info, result)
def _intercept_server_stream(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) as span: inject(mutable_metadata, setter=_carrier_setter) metadata = tuple(mutable_metadata.items()) rpc_info = RpcInfo( full_method=client_info.full_method, metadata=metadata, timeout=client_info.timeout, ) if client_info.is_client_stream: rpc_info.request = request_or_iterator try: yield from invoker(request_or_iterator, metadata) except grpc.RpcError as err: span.set_status(Status(StatusCode.ERROR)) span.set_attribute(SpanAttributes.RPC_GRPC_STATUS_CODE, err.code().value[0]) raise err
def _intercept_server_stream(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) as span: _inject_span_context(mutable_metadata) metadata = tuple(mutable_metadata.items()) rpc_info = RpcInfo( full_method=client_info.full_method, metadata=metadata, timeout=client_info.timeout, ) if client_info.is_client_stream: rpc_info.request = request_or_iterator try: result = invoker(request_or_iterator, metadata) for response in result: yield response except grpc.RpcError as err: span.set_status(Status(StatusCode.ERROR)) span.set_attribute("rpc.grpc.status_code", err.code().value[0]) raise err
def intercept_unary(self, request, metadata, client_info, invoker): if not metadata: mutable_metadata = OrderedDict() else: mutable_metadata = OrderedDict(metadata) with self._start_guarded_span(client_info.full_method) as guarded_span: with self._metrics_recorder.record_latency( client_info.full_method): _inject_span_context(mutable_metadata) metadata = tuple(mutable_metadata.items()) # If protobuf is used, we can record the bytes in/out. Otherwise, we have no way # to get the size of the request/response properly, so don't record anything if "ByteSize" in dir(request): self._metrics_recorder.record_bytes_out( request.ByteSize(), client_info.full_method) rpc_info = RpcInfo( full_method=client_info.full_method, metadata=metadata, timeout=client_info.timeout, request=request, ) try: result = invoker(request, metadata) except grpc.RpcError: guarded_span.generated_span.set_status( Status(StatusCode.ERROR)) raise return self._trace_result(guarded_span, rpc_info, result, client_info)
def intercept_unary(self, request, metadata, client_info, invoker): if not metadata: mutable_metadata = OrderedDict() else: mutable_metadata = OrderedDict(metadata) with self._start_guarded_span(client_info.full_method) as guarded_span: inject(mutable_metadata, setter=_carrier_setter) metadata = tuple(mutable_metadata.items()) rpc_info = RpcInfo( full_method=client_info.full_method, metadata=metadata, timeout=client_info.timeout, request=request, ) try: result = invoker(request, metadata) except grpc.RpcError as err: guarded_span.generated_span.set_status(Status( StatusCode.ERROR)) guarded_span.generated_span.set_attribute( SpanAttributes.RPC_GRPC_STATUS_CODE, err.code().value[0]) raise err return self._trace_result(guarded_span, rpc_info, result)
def _intercept_server_stream(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) as span: with self._metrics_recorder.record_latency( client_info.full_method): _inject_span_context(mutable_metadata) metadata = tuple(mutable_metadata.items()) rpc_info = RpcInfo( full_method=client_info.full_method, metadata=metadata, timeout=client_info.timeout, ) if client_info.is_client_stream: rpc_info.request = request_or_iterator request_or_iterator = self._bytes_out_iterator_wrapper( request_or_iterator, client_info) else: if "ByteSize" in dir(request_or_iterator): self._metrics_recorder.record_bytes_out( request_or_iterator.ByteSize(), client_info.full_method, ) try: result = invoker(request_or_iterator, metadata) # Rewrap the result stream into a generator, and record the bytes received for response in result: if "ByteSize" in dir(response): self._metrics_recorder.record_bytes_in( response.ByteSize(), client_info.full_method) yield response except grpc.RpcError: span.set_status(Status(StatusCode.ERROR)) raise
def intercept_stream(self, request_or_iterator, metadata, client_info, invoker): if client_info.is_server_stream: return self._intercept_server_stream(request_or_iterator, metadata, client_info, invoker) if not metadata: mutable_metadata = OrderedDict() else: mutable_metadata = OrderedDict(metadata) with self._start_guarded_span(client_info.full_method) as guarded_span: with self._metrics_recorder.record_latency( client_info.full_method): _inject_span_context(mutable_metadata) metadata = tuple(mutable_metadata.items()) rpc_info = RpcInfo( full_method=client_info.full_method, metadata=metadata, timeout=client_info.timeout, request=request_or_iterator, ) rpc_info.request = request_or_iterator request_or_iterator = self._bytes_out_iterator_wrapper( request_or_iterator, client_info) try: result = invoker(request_or_iterator, metadata) except grpc.RpcError: guarded_span.generated_span.set_status( Status(StatusCode.ERROR)) raise return self._trace_result(guarded_span, rpc_info, result, client_info)
def _intercept(self, request, metadata, client_info, invoker): if context.get_value(_SUPPRESS_INSTRUMENTATION_KEY): return invoker(request, metadata) if not metadata: mutable_metadata = OrderedDict() 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, ) as span: result = None try: inject(mutable_metadata, setter=_carrier_setter) metadata = tuple(mutable_metadata.items()) rpc_info = RpcInfo( full_method=client_info.full_method, metadata=metadata, timeout=client_info.timeout, request=request, ) result = invoker(request, metadata) except Exception as exc: if isinstance(exc, grpc.RpcError): span.set_attribute( SpanAttributes.RPC_GRPC_STATUS_CODE, exc.code().value[0], ) span.set_status( Status( status_code=StatusCode.ERROR, description=f"{type(exc).__name__}: {exc}", )) span.record_exception(exc) raise exc finally: if not result: span.end() return self._trace_result(span, rpc_info, result)