fix(boto3): Finish StreamingBody span correctly

sentry_sdk/integrations/boto3/_instrumentation.py

@@ -1,10 +1,10 @@
 from typing import TYPE_CHECKING
 
 import sentry_sdk
-from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS
 from sentry_sdk.integrations import DidNotEnable
 from sentry_sdk.integrations.boto3.consts import ORIGIN
-from sentry_sdk.traces import StreamedSpan
+from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan
 from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span
 from sentry_sdk.tracing_utils import (
     add_http_breadcrumb,
@@ -164,26 +164,44 @@ def _replace_header(request: "AWSRequest", key: str, value: str) -> None:
             )
 
 
-def _sentry_after_call(
-    context: "Dict[str, Any]", parsed: "Dict[str, Any]", **kwargs: "Any"
+def _finish_span(
+    span: "Union[Span, StreamedSpan]",
+    error: "Optional[BaseException]" = None,
 ) -> None:
-    span: "Optional[Union[Span, StreamedSpan]]" = context.pop("_sentrysdk_span", None)
+    with capture_internal_exceptions():
+        if not isinstance(span, StreamedSpan):
+            if error is not None:
+                span.set_status(SPANSTATUS.INTERNAL_ERROR)
+            span.finish()
+            return
 
-    # Span could be absent if the integration is disabled.
-    if span is None:
-        return
+        if error is None:
+            span.end()
+        else:
+            span.__exit__(type(error), error, error.__traceback__)
 
-    span.__exit__(None, None, None)
+
+def _instrument_streaming_body(
+    span: "Union[Span, StreamedSpan]", parsed: "Dict[str, Any]"
+) -> bool:
+    if isinstance(span, NoOpStreamedSpan):
+        return False
 
     body = parsed.get("Body")
     if not isinstance(body, StreamingBody):
-        return
+        return False
 
     streaming_span: "Union[Span, StreamedSpan]"
     if isinstance(span, StreamedSpan):
         streaming_span = sentry_sdk.traces.start_span(
             name=span.name,
+            # `parent_span` is set explicitly to the boto span.
             parent_span=span,
+            # avoid making the streaming span the current span on the scope since the application might
+            # keep `StreamingBody` open before reading it. Otherwise: 1. when the streamingspan ends it
+            # could restore the parent span on the scope, breaking the parent-child relation of newly
+            # created spans; 2. newly created spans would be attached to the streaming span.
+            active=False,
             attributes={
                 "sentry.op": OP.HTTP_CLIENT_STREAM,
                 "sentry.origin": ORIGIN,
@@ -198,35 +216,92 @@ def _sentry_after_call(
 
     orig_read = body.read
     orig_close = body.close
+    raw_stream = body._raw_stream  # type: ignore[attr-defined]
+    orig_raw_close = raw_stream.close
+    finished = False
+    read_in_progress = False
+
+    def finish_span(error: "Optional[BaseException]" = None) -> None:
+        nonlocal finished
+        if finished:
+            return
+
+        finished = True
+        _finish_span(streaming_span, error)
+
+    def content_length_reached() -> bool:
+        content_length = getattr(body, "_content_length", None)
+        amount_read = getattr(body, "_amount_read", None)
+        return (
+            content_length is not None
+            and amount_read is not None
+            and amount_read >= int(content_length)
+        )
 
     def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes:
+        nonlocal read_in_progress
+        read_in_progress = True
         try:
-            ret = orig_read(*args, **kwargs)
-            if ret:
-                return ret
-
-            if isinstance(streaming_span, StreamedSpan):
-                streaming_span.end()
-            else:
-                streaming_span.finish()
-            return ret
-        except Exception:
-            if isinstance(streaming_span, StreamedSpan):
-                streaming_span.end()
-            else:
-                streaming_span.finish()
+            read_return_value = orig_read(*args, **kwargs)
+            with capture_internal_exceptions():
+                amount_of_bytes_requested = args[0] if args else kwargs.get("amt")
+                if (
+                    amount_of_bytes_requested is None
+                    or amount_of_bytes_requested < 0
+                    or (amount_of_bytes_requested > 0 and not read_return_value)
+                    or content_length_reached()
+                ):
+                    finish_span()
+            return read_return_value
+        except BaseException as error:
+            finish_span(error)
             raise
-
-    body.read = sentry_streaming_body_read  # type: ignore
+        finally:
+            read_in_progress = False
 
     def sentry_streaming_body_close(*args: "Any", **kwargs: "Any") -> None:
-        if isinstance(streaming_span, StreamedSpan):
-            streaming_span.end()
-        else:
-            streaming_span.finish()
-        orig_close(*args, **kwargs)
+        try:
+            orig_close(*args, **kwargs)
+            finish_span()
+        except BaseException as error:
+            finish_span(error)
+            raise
+
+    def sentry_raw_stream_close(*args: "Any", **kwargs: "Any") -> None:
+        try:
+            orig_raw_close(*args, **kwargs)
+            if not read_in_progress:
+                finish_span()
+        except BaseException as error:
+            finish_span(error)
+            raise
+
+    try:
+        # StreamingBody.__exit__ closes `_raw_stream` directly, bypassing
+        # StreamingBody.close(), so both levels need to be instrumented.
+        raw_stream.close = sentry_raw_stream_close
+        body.read = sentry_streaming_body_read  # type: ignore
+        body.close = sentry_streaming_body_close  # type: ignore
+    except Exception:
+        finish_span()
+        raise
 
-    body.close = sentry_streaming_body_close  # type: ignore
+    return True
+
+
+def _sentry_after_call(
+    context: "Dict[str, Any]", parsed: "Dict[str, Any]", **kwargs: "Any"
+) -> None:
+    span: "Optional[Union[Span, StreamedSpan]]" = context.pop("_sentrysdk_span", None)
+
+    # Span could be absent if the integration is disabled.
+    if span is None:
+        return
+
+    span.__exit__(None, None, None)
+
+    with capture_internal_exceptions():
+        _instrument_streaming_body(span, parsed)
 
 
 def _sentry_after_call_error(

tests/integrations/boto3/test_client.py

@@ -1,6 +1,134 @@
+import boto3
+import pytest
+from botocore.awsrequest import AWSResponse
+from botocore.config import Config
+
+import sentry_sdk
+from sentry_sdk.consts import OP
 from sentry_sdk.integrations.boto3 import Boto3Integration
+from tests.integrations.boto3.aws_mock import Body
+
+session = boto3.Session(  # type: ignore[attr-defined]
+    aws_access_key_id="-",
+    aws_secret_access_key="-",
+    region_name="eu-north-1",
+)
 
 
 def test_public_api():
     assert Boto3Integration.__module__ == "sentry_sdk.integrations.boto3"
     assert Boto3Integration.identifier == "boto3"
+
+
+@pytest.fixture
+def client_factory(sentry_init, monkeypatch, span_streaming):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[Boto3Integration()],
+        trace_lifecycle="stream" if span_streaming else "static",
+        # avoid SDK's machine hostname being used as server name.
+        server_name="",
+    )
+    # remove retry delay to speed up tests
+    monkeypatch.setattr("botocore.endpoint.time.sleep", lambda delay: None)
+
+    def make_client(service_name="s3", attempt_count=1, **client_kwargs):
+        return session.client(
+            service_name,
+            config=Config(
+                # `total_max_attempts` includes the initial request.
+                retries={"total_max_attempts": attempt_count, "mode": "standard"}
+            ),
+            **client_kwargs,
+        )
+
+    return make_client
+
+
+def _capture_boto3_spans_by_op(invoke_client_method, capture_items, span_streaming):
+    items = capture_items()
+
+    if span_streaming:
+        with sentry_sdk.traces.start_span(name="parent"):  # type: ignore[attr-defined]
+            invoke_client_method()
+
+        sentry_sdk.flush()
+        spans = [
+            item.payload
+            for item in items
+            if item.type == "span"
+            and item.payload["attributes"].get("sentry.origin")
+            == Boto3Integration.origin
+        ]
+    else:
+        with sentry_sdk.start_transaction():
+            invoke_client_method()
+
+        transaction = next(item.payload for item in items if item.type == "transaction")
+        spans = [
+            span
+            for span in transaction["spans"]
+            if span["origin"] == Boto3Integration.origin
+        ]
+
+    spans_by_op = {}
+    for span in spans:
+        op = span["attributes"].get("sentry.op") if span_streaming else span["op"]
+        spans_by_op.setdefault(op, []).append(span)
+    return spans_by_op
+
+
+def _assert_span_finished(span, span_streaming):
+    finished_timestamp = "end_timestamp" if span_streaming else "timestamp"
+    assert span[finished_timestamp] is not None
+
+
+def _assert_one_failed_span(spans, span_streaming):
+    assert len(spans) == 1
+    assert spans[0]["status"] in ("error", "internal_error")
+    _assert_span_finished(spans[0], span_streaming)
+
+
+@pytest.mark.parametrize("span_streaming", [True, False])
+def test_streaming_body_read_failure_finishes_stream_span(
+    capture_items,
+    client_factory,
+    span_streaming,
+):
+    client = client_factory()
+    original_exception = OSError("stream read failed")
+
+    class _FailingBody(Body):
+        def __init__(self, exception):
+            super().__init__(b"")
+            self._exception = exception
+
+        def read(self, *args, **kwargs):
+            # urllib3 closes the response before propagating some read failures.
+            self.close()
+            raise self._exception
+
+    def respond(request, **kwargs):
+        return AWSResponse(
+            request.url,
+            200,
+            {"content-length": "1"},
+            _FailingBody(original_exception),
+        )
+
+    client.meta.events.register("before-send", respond)
+
+    def invoke_client_method_and_read_body():
+        body = client.get_object(Bucket="bucket", Key="foo")["Body"]
+        with pytest.raises(OSError) as exc_info:
+            body.read()
+        assert exc_info.value is original_exception
+
+    spans_by_op = _capture_boto3_spans_by_op(
+        invoke_client_method_and_read_body, capture_items, span_streaming
+    )
+    client_spans = spans_by_op.get(OP.HTTP_CLIENT, [])
+    stream_spans = spans_by_op.get(OP.HTTP_CLIENT_STREAM, [])
+
+    assert len(client_spans) == 1
+    _assert_one_failed_span(stream_spans, span_streaming)