From 7d08e57b76f905150abd5fe87f53d6e68f37e9ea Mon Sep 17 00:00:00 2001 From: David Montague <35119617+dmontagu@users.noreply.github.com> Date: Fri, 20 Dec 2024 15:08:53 -0700 Subject: [PATCH 1/3] Remove double record exception --- logfire/_internal/main.py | 48 +++++++++---------------------------- logfire/_internal/tracer.py | 6 +++++ tests/test_logfire.py | 22 +++++++++++++++++ 3 files changed, 39 insertions(+), 37 deletions(-) diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index d3c17752d..2a9ac5d4c 100644 --- a/logfire/_internal/main.py +++ b/logfire/_internal/main.py @@ -1948,13 +1948,13 @@ def __init__(self, span: trace_api.Span) -> None: self._token = context_api.attach(trace_api.set_span_in_context(self._span)) def __enter__(self) -> FastLogfireSpan: + self._span.__enter__() return self @handle_internal_errors() def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any) -> None: context_api.detach(self._token) - _exit_span(self._span, exc_value) - self._span.end() + self._span.__exit__(exc_type, exc_value, traceback) # Changes to this class may need to be reflected in `FastLogfireSpan` and `NoopSpan` as well. @@ -1990,6 +1990,7 @@ def __enter__(self) -> LogfireSpan: attributes=self._otlp_attributes, links=self._links, ) + self._span.__enter__() if self._token is None: # pragma: no branch self._token = context_api.attach(trace_api.set_span_in_context(self._span)) @@ -1999,14 +2000,17 @@ def __enter__(self) -> LogfireSpan: def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any) -> None: if self._token is None: # pragma: no cover return + assert self._span is not None context_api.detach(self._token) self._token = None - - assert self._span is not None - _exit_span(self._span, exc_value) - - self.end() + if self._span.is_recording(): + with handle_internal_errors(): + if self._added_attributes: + self._span.set_attribute( + ATTRIBUTES_JSON_SCHEMA_KEY, attributes_json_schema(self._json_schema_properties) + ) + self._span.__exit__(exc_type, exc_value, traceback) @property def message_template(self) -> str | None: # pragma: no cover @@ -2032,26 +2036,6 @@ def message(self) -> str: def message(self, message: str): self._set_attribute(ATTRIBUTES_MESSAGE_KEY, message) - def end(self, end_time: int | None = None) -> None: - """Sets the current time as the span's end time. - - The span's end time is the wall time at which the operation finished. - - Only the first call to this method is recorded, further calls are ignored so you - can call this within the span's context manager to end it before the context manager - exits. - """ - if self._span is None: # pragma: no cover - raise RuntimeError('Span has not been started') - if self._span.is_recording(): - with handle_internal_errors(): - if self._added_attributes: - self._span.set_attribute( - ATTRIBUTES_JSON_SCHEMA_KEY, attributes_json_schema(self._json_schema_properties) - ) - - self._span.end(end_time) - @handle_internal_errors() def set_attribute(self, key: str, value: Any) -> None: """Sets an attribute on the span. @@ -2183,16 +2167,6 @@ def is_recording(self) -> bool: return False -def _exit_span(span: trace_api.Span, exception: BaseException | None) -> None: - if not span.is_recording(): - return - - # record exception if present - # isinstance is to ignore BaseException - if isinstance(exception, Exception): - record_exception(span, exception, escaped=True) - - AttributesValueType = TypeVar('AttributesValueType', bound=Union[Any, otel_types.AttributeValue]) diff --git a/logfire/_internal/tracer.py b/logfire/_internal/tracer.py index 5d7ad8410..b2206c8a2 100644 --- a/logfire/_internal/tracer.py +++ b/logfire/_internal/tracer.py @@ -171,6 +171,12 @@ def record_exception( timestamp = timestamp or self.ns_timestamp_generator() record_exception(self.span, exception, attributes=attributes, timestamp=timestamp, escaped=escaped) + def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any) -> None: + if self.is_recording() and isinstance(exc_value, BaseException): + self.record_exception(exc_value, escaped=True) + + self.end() + if not TYPE_CHECKING: # pragma: no branch # for ReadableSpan def __getattr__(self, name: str) -> Any: diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 420e8d399..eb7493bc2 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -9,6 +9,7 @@ from functools import partial from logging import getLogger from typing import Any, Callable +from unittest.mock import patch import pytest from dirty_equals import IsInt, IsJson, IsStr @@ -36,6 +37,7 @@ ) from logfire._internal.formatter import FormattingFailedWarning, InspectArgumentsFailedWarning from logfire._internal.main import NoopSpan +from logfire._internal.tracer import record_exception from logfire._internal.utils import is_instrumentation_suppressed from logfire.integrations.logging import LogfireLoggingHandler from logfire.testing import TestExporter @@ -3171,3 +3173,23 @@ def test_suppress_scopes(exporter: TestExporter, metrics_reader: InMemoryMetricR } ] ) + + +def test_logfire_span_records_exceptions_once(): + # logfire.configure(send_to_logfire=False, console=False, additional_span_processors=[BatchSpanProcessor(exporter)]) + n_calls_to_record_exception = 0 + + def patched_record_exception(*args: Any, **kwargs: Any) -> Any: + nonlocal n_calls_to_record_exception + n_calls_to_record_exception += 1 + + return record_exception(*args, **kwargs) + + with patch('logfire._internal.tracer.record_exception', patched_record_exception), patch( + 'logfire._internal.main.record_exception', patched_record_exception + ): + with pytest.raises(RuntimeError): + with logfire.span('foo'): + raise RuntimeError('error') + + assert n_calls_to_record_exception == 1 From 83a60a9ebc7843a93335dcf1b5ffe8398a4f83a3 Mon Sep 17 00:00:00 2001 From: David Montague <35119617+dmontagu@users.noreply.github.com> Date: Fri, 20 Dec 2024 15:18:44 -0700 Subject: [PATCH 2/3] Fix tests --- tests/test_logfire.py | 25 ++++++++++++------------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index eb7493bc2..30b4ffdbb 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -1362,8 +1362,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'root (pending)', 'context': {'trace_id': 1, 'span_id': 2, 'is_remote': False}, 'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, - 'start_time': 5000000000, - 'end_time': 5000000000, + 'start_time': 7000000000, + 'end_time': 7000000000, 'attributes': { 'code.filepath': 'test_logfire.py', 'code.lineno': 123, @@ -1378,8 +1378,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'child (pending)', 'context': {'trace_id': 1, 'span_id': 4, 'is_remote': False}, 'parent': {'trace_id': 1, 'span_id': 3, 'is_remote': False}, - 'start_time': 6000000000, - 'end_time': 6000000000, + 'start_time': 8000000000, + 'end_time': 8000000000, 'attributes': { 'code.filepath': 'test_logfire.py', 'code.lineno': 123, @@ -1394,8 +1394,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'test1', 'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False}, 'parent': {'trace_id': 1, 'span_id': 3, 'is_remote': False}, - 'start_time': 7000000000, - 'end_time': 7000000000, + 'start_time': 9000000000, + 'end_time': 9000000000, 'attributes': { 'logfire.span_type': 'log', 'logfire.level_num': 9, @@ -1410,8 +1410,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'test2', 'context': {'trace_id': 1, 'span_id': 6, 'is_remote': False}, 'parent': {'trace_id': 1, 'span_id': 3, 'is_remote': False}, - 'start_time': 8000000000, - 'end_time': 8000000000, + 'start_time': 10000000000, + 'end_time': 10000000000, 'attributes': { 'logfire.span_type': 'log', 'logfire.level_num': 9, @@ -1427,8 +1427,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'child', 'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False}, 'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, - 'start_time': 6000000000, - 'end_time': 9000000000, + 'start_time': 8000000000, + 'end_time': 11000000000, 'attributes': { 'code.filepath': 'test_logfire.py', 'code.lineno': 123, @@ -1442,8 +1442,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'root', 'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, 'parent': None, - 'start_time': 5000000000, - 'end_time': 10000000000, + 'start_time': 7000000000, + 'end_time': 12000000000, 'attributes': { 'code.filepath': 'test_logfire.py', 'code.lineno': 123, @@ -3176,7 +3176,6 @@ def test_suppress_scopes(exporter: TestExporter, metrics_reader: InMemoryMetricR def test_logfire_span_records_exceptions_once(): - # logfire.configure(send_to_logfire=False, console=False, additional_span_processors=[BatchSpanProcessor(exporter)]) n_calls_to_record_exception = 0 def patched_record_exception(*args: Any, **kwargs: Any) -> Any: From fe7d66bb123962ab91d6a4b95f9f5a6b1b276c4b Mon Sep 17 00:00:00 2001 From: David Montague <35119617+dmontagu@users.noreply.github.com> Date: Fri, 20 Dec 2024 15:38:07 -0700 Subject: [PATCH 3/3] Revert tests to old behavior --- logfire/_internal/tracer.py | 8 ++++---- tests/test_logfire.py | 24 ++++++++++++------------ 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/logfire/_internal/tracer.py b/logfire/_internal/tracer.py index b2206c8a2..1b8b9fc4a 100644 --- a/logfire/_internal/tracer.py +++ b/logfire/_internal/tracer.py @@ -172,10 +172,10 @@ def record_exception( record_exception(self.span, exception, attributes=attributes, timestamp=timestamp, escaped=escaped) def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any) -> None: - if self.is_recording() and isinstance(exc_value, BaseException): - self.record_exception(exc_value, escaped=True) - - self.end() + if self.is_recording(): + if isinstance(exc_value, BaseException): + self.record_exception(exc_value, escaped=True) + self.end() if not TYPE_CHECKING: # pragma: no branch # for ReadableSpan diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 30b4ffdbb..be83dc123 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -1362,8 +1362,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'root (pending)', 'context': {'trace_id': 1, 'span_id': 2, 'is_remote': False}, 'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, - 'start_time': 7000000000, - 'end_time': 7000000000, + 'start_time': 5000000000, + 'end_time': 5000000000, 'attributes': { 'code.filepath': 'test_logfire.py', 'code.lineno': 123, @@ -1378,8 +1378,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'child (pending)', 'context': {'trace_id': 1, 'span_id': 4, 'is_remote': False}, 'parent': {'trace_id': 1, 'span_id': 3, 'is_remote': False}, - 'start_time': 8000000000, - 'end_time': 8000000000, + 'start_time': 6000000000, + 'end_time': 6000000000, 'attributes': { 'code.filepath': 'test_logfire.py', 'code.lineno': 123, @@ -1394,8 +1394,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'test1', 'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False}, 'parent': {'trace_id': 1, 'span_id': 3, 'is_remote': False}, - 'start_time': 9000000000, - 'end_time': 9000000000, + 'start_time': 7000000000, + 'end_time': 7000000000, 'attributes': { 'logfire.span_type': 'log', 'logfire.level_num': 9, @@ -1410,8 +1410,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'test2', 'context': {'trace_id': 1, 'span_id': 6, 'is_remote': False}, 'parent': {'trace_id': 1, 'span_id': 3, 'is_remote': False}, - 'start_time': 10000000000, - 'end_time': 10000000000, + 'start_time': 8000000000, + 'end_time': 8000000000, 'attributes': { 'logfire.span_type': 'log', 'logfire.level_num': 9, @@ -1427,8 +1427,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'child', 'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False}, 'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, - 'start_time': 8000000000, - 'end_time': 11000000000, + 'start_time': 6000000000, + 'end_time': 9000000000, 'attributes': { 'code.filepath': 'test_logfire.py', 'code.lineno': 123, @@ -1442,8 +1442,8 @@ def test_logfire_with_its_own_config(exporter: TestExporter, config_kwargs: dict 'name': 'root', 'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, 'parent': None, - 'start_time': 7000000000, - 'end_time': 12000000000, + 'start_time': 5000000000, + 'end_time': 10000000000, 'attributes': { 'code.filepath': 'test_logfire.py', 'code.lineno': 123,