Skip to content

Commit

Permalink
Use is_recording flag in asgi, pyramid, aiohttp instrumentation (#1142)
Browse files Browse the repository at this point in the history
  • Loading branch information
lzchen committed Sep 22, 2020
1 parent 65fcb3a commit 7718d88
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 15 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -156,38 +156,41 @@ async def __call__(self, scope, receive, send):
propagators.extract(get_header_from_scope, scope)
)
span_name, additional_attributes = self.span_details_callback(scope)
attributes = collect_request_attributes(scope)
attributes.update(additional_attributes)

try:
with self.tracer.start_as_current_span(
span_name + " asgi",
kind=trace.SpanKind.SERVER,
attributes=attributes,
):
span_name + " asgi", kind=trace.SpanKind.SERVER,
) as span:
if span.is_recording():
attributes = collect_request_attributes(scope)
attributes.update(additional_attributes)
for key, value in attributes.items():
span.set_attribute(key, value)

@wraps(receive)
async def wrapped_receive():
with self.tracer.start_as_current_span(
span_name + " asgi." + scope["type"] + ".receive"
) as receive_span:
message = await receive()
if message["type"] == "websocket.receive":
set_status_code(receive_span, 200)
receive_span.set_attribute("type", message["type"])
if receive_span.is_recording():
if message["type"] == "websocket.receive":
set_status_code(receive_span, 200)
receive_span.set_attribute("type", message["type"])
return message

@wraps(send)
async def wrapped_send(message):
with self.tracer.start_as_current_span(
span_name + " asgi." + scope["type"] + ".send"
) as send_span:
if message["type"] == "http.response.start":
status_code = message["status"]
set_status_code(send_span, status_code)
elif message["type"] == "websocket.send":
set_status_code(send_span, 200)
send_span.set_attribute("type", message["type"])
if send_span.is_recording():
if message["type"] == "http.response.start":
status_code = message["status"]
set_status_code(send_span, status_code)
elif message["type"] == "websocket.send":
set_status_code(send_span, 200)
send_span.set_attribute("type", message["type"])
await send(message)

await self.app(scope, wrapped_receive, wrapped_send)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,23 @@ def test_basic_asgi_call(self):
outputs = self.get_all_output()
self.validate_outputs(outputs)

def test_wsgi_not_recording(self):
mock_tracer = mock.Mock()
mock_span = mock.Mock()
mock_span.is_recording.return_value = False
mock_tracer.start_as_current_span.return_value = mock_span
mock_tracer.start_as_current_span.return_value.__enter__ = mock_span
mock_tracer.start_as_current_span.return_value.__exit__ = mock_span
with mock.patch("opentelemetry.trace.get_tracer") as tracer:
tracer.return_value = mock_tracer
app = otel_asgi.OpenTelemetryMiddleware(simple_asgi)
self.seed_app(app)
self.send_default_request()
self.assertFalse(mock_span.is_recording())
self.assertTrue(mock_span.is_recording.called)
self.assertFalse(mock_span.set_attribute.called)
self.assertFalse(mock_span.set_status.called)

def test_asgi_exc_info(self):
"""Test that exception information is emitted as expected."""
app = otel_asgi.OpenTelemetryMiddleware(error_asgi)
Expand Down

0 comments on commit 7718d88

Please sign in to comment.