fix: address #4681 - Django Instrumentation: Spans close prematurely for StreamingHttpResponse

Closes #4681

Author: Mergepath

instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware/otel_middleware.py

@@ -259,7 +259,7 @@ def process_request(self, request):
             for key, value in attributes.items():
                 span.set_attribute(key, value)
 
-        activation = use_span(span, end_on_exit=True)
+        activation = use_span(span, end_on_exit=False)
         activation.__enter__()  # pylint: disable=unnecessary-dunder-call
         request_start_time = default_timer()
         request.META[self._environ_timer_key] = request_start_time
@@ -394,29 +394,40 @@ def process_response(self, request, response):
                 except Exception:  # pylint: disable=broad-exception-caught
                     _logger.exception("Exception raised by response_hook")
 
-        if request_start_time is not None:
-            duration_s = default_timer() - request_start_time
-            if self._duration_histogram_old:
-                duration_attrs_old = _parse_duration_attrs(
-                    duration_attrs, _StabilityMode.DEFAULT
-                )
-                # http.target to be included in old semantic conventions
-                target = duration_attrs.get(HTTP_TARGET)
-                if target:
-                    duration_attrs_old[HTTP_TARGET] = target
-                self._duration_histogram_old.record(
-                    max(round(duration_s * 1000), 0),
-                    duration_attrs_old,
-                )
-            if self._duration_histogram_new:
-                duration_attrs_new = _parse_duration_attrs(
-                    duration_attrs, _StabilityMode.HTTP
-                )
-                self._duration_histogram_new.record(
-                    max(duration_s, 0),
-                    duration_attrs_new,
-                )
-        self._active_request_counter.add(-1, active_requests_count_attrs)
+        finalized = False
+
+        def finalize_response():
+            nonlocal finalized
+            if finalized:
+                return
+            finalized = True
+
+            if request_start_time is not None:
+                duration_s = default_timer() - request_start_time
+                if self._duration_histogram_old:
+                    duration_attrs_old = _parse_duration_attrs(
+                        duration_attrs, _StabilityMode.DEFAULT
+                    )
+                    # http.target to be included in old semantic conventions
+                    target = duration_attrs.get(HTTP_TARGET)
+                    if target:
+                        duration_attrs_old[HTTP_TARGET] = target
+                    self._duration_histogram_old.record(
+                        max(round(duration_s * 1000), 0),
+                        duration_attrs_old,
+                    )
+                if self._duration_histogram_new:
+                    duration_attrs_new = _parse_duration_attrs(
+                        duration_attrs, _StabilityMode.HTTP
+                    )
+                    self._duration_histogram_new.record(
+                        max(duration_s, 0),
+                        duration_attrs_new,
+                    )
+            self._active_request_counter.add(-1, active_requests_count_attrs)
+
+            if span:
+                span.end()
 
         if activation and span:
             if exception:
@@ -432,6 +443,16 @@ def process_response(self, request, response):
             detach(request.META.get(self._environ_token))
             request.META.pop(self._environ_token)
 
+        original_close = response.close
+
+        def close():
+            try:
+                return original_close()
+            finally:
+                finalize_response()
+
+        response.close = close
+
         return response
 
 

instrumentation/opentelemetry-instrumentation-django/tests/test_middleware.py

@@ -58,6 +58,7 @@
     excluded_noarg2,
     response_with_custom_header,
     route_span_name,
+    streaming,
     traced,
     traced_template,
 )
@@ -77,6 +78,7 @@ def path(path_argument, *args, **kwargs):
 
 urlpatterns = [
     re_path(r"^traced/", traced),
+    re_path(r"^streaming/", streaming),
     re_path(r"^traced_custom_header/", response_with_custom_header),
     re_path(r"^route/(?P<year>[0-9]{4})/template/$", traced_template),
     re_path(r"^error/", error),
@@ -242,6 +244,39 @@ def test_traced_get(self):
         self.assertEqual(span.attributes["http.scheme"], "http")
         self.assertEqual(span.attributes["http.status_code"], 200)
 
+    def test_streaming_response_span_closes_after_response_close(self):
+        response = Client().get("/streaming/")
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 0)
+
+        streaming_content = iter(response.streaming_content)
+        self.assertEqual(next(streaming_content), b"streaming ")
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 0)
+
+        response.close()
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+
+        response.close()
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+
+        span = spans[0]
+        self.assertEqual(span.name, "GET ^streaming/" if DJANGO_2_2 else "GET")
+        self.assertEqual(span.kind, SpanKind.SERVER)
+        self.assertEqual(span.status.status_code, StatusCode.UNSET)
+        self.assertEqual(span.attributes["http.method"], "GET")
+        self.assertEqual(
+            span.attributes["http.url"],
+            "http://testserver/streaming/",
+        )
+        if DJANGO_2_2:
+            self.assertEqual(span.attributes["http.route"], "^streaming/")
+        self.assertEqual(span.attributes["http.scheme"], "http")
+        self.assertEqual(span.attributes["http.status_code"], 200)
+
     def test_traced_get_new_semconv(self):
         Client().get("/traced/")
 

instrumentation/opentelemetry-instrumentation-django/tests/test_middleware_asgi.py

@@ -72,6 +72,7 @@
     async_excluded_noarg,
     async_excluded_noarg2,
     async_route_span_name,
+    async_streaming,
     async_traced,
     async_traced_template,
     async_with_custom_header,
@@ -87,6 +88,7 @@
 
 urlpatterns = [
     re_path(r"^traced/", async_traced),
+    re_path(r"^streaming/", async_streaming),
     re_path(r"^traced_custom_header/", async_with_custom_header),
     re_path(r"^route/(?P<year>[0-9]{4})/template/$", async_traced_template),
     re_path(r"^error/", async_error),
@@ -248,6 +250,44 @@ async def test_traced_get(self):
         self.assertEqual(span.attributes[HTTP_SCHEME], "http")
         self.assertEqual(span.attributes[HTTP_STATUS_CODE], 200)
 
+    async def test_streaming_response_span_closes_after_response_close(self):
+        response = await self.async_client.get("/streaming/")
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 0)
+
+        if hasattr(response.streaming_content, "__aiter__"):
+            streaming_content = response.streaming_content.__aiter__()
+            self.assertEqual(
+                await streaming_content.__anext__(), b"streaming "
+            )
+        else:
+            streaming_content = iter(response.streaming_content)
+            self.assertEqual(next(streaming_content), b"streaming ")
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 0)
+
+        response.close()
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+
+        response.close()
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+
+        span = spans[0]
+        self.assertEqual(span.name, "GET ^streaming/")
+        self.assertEqual(span.kind, SpanKind.SERVER)
+        self.assertEqual(span.status.status_code, StatusCode.UNSET)
+        self.assertEqual(span.attributes[HTTP_METHOD], "GET")
+        self.assertEqual(
+            span.attributes[HTTP_URL],
+            "http://testserver/streaming/",
+        )
+        self.assertEqual(span.attributes[HTTP_ROUTE], "^streaming/")
+        self.assertEqual(span.attributes[HTTP_SCHEME], "http")
+        self.assertEqual(span.attributes[HTTP_STATUS_CODE], 200)
+
     async def test_traced_get_new_semconv(self):
         await self.async_client.get("/traced/")
 

instrumentation/opentelemetry-instrumentation-django/tests/views.py

@@ -1,13 +1,18 @@
 # Copyright The OpenTelemetry Authors
 # SPDX-License-Identifier: Apache-2.0
 
-from django.http import HttpResponse
+from django import VERSION
+from django.http import HttpResponse, StreamingHttpResponse
 
 
 def traced(request):  # pylint: disable=unused-argument
     return HttpResponse()
 
 
+def streaming(request):  # pylint: disable=unused-argument
+    return StreamingHttpResponse(iter([b"streaming ", b"response"]))
+
+
 def traced_template(request, year):  # pylint: disable=unused-argument
     return HttpResponse()
 
@@ -50,6 +55,17 @@ async def async_traced(request):  # pylint: disable=unused-argument
     return HttpResponse()
 
 
+async def async_streaming(request):  # pylint: disable=unused-argument
+    if VERSION < (4, 2):
+        return StreamingHttpResponse(iter([b"streaming ", b"response"]))
+
+    async def streaming_content():
+        yield b"streaming "
+        yield b"response"
+
+    return StreamingHttpResponse(streaming_content())
+
+
 async def async_traced_template(request, year):  # pylint: disable=unused-argument
     return HttpResponse()