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)