chore: compatiable flask 3.1+

Author: fatelei

CHANGELOG.md

@@ -43,7 +43,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `opentelemetry-instrumentation-aiohttp-server`: delay initialization of tracer, meter and excluded urls to instrumentation for testability
   ([#3836](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3836))
 - `opentelemetry-instrumentation-elasticsearch`: Enhance elasticsearch query body sanitization
-  ([#3919](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3919)) 
+  ([#3919](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3919))
+- `opentelemetry-instrumentation-flask`: Add Flask 3.1+ compatibility with proper context cleanup for streaming responses to prevent memory leaks and token reuse
+  ([#3937](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3938)) 
 
 
 ## Version 1.38.0/0.59b0 (2025-10-16)

instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py

@@ -309,7 +309,16 @@ def response_hook(span: Span, status: str, response_headers: List):
 
 flask_version = version("flask")
 
-if package_version.parse(flask_version) >= package_version.parse("2.2.0"):
+if package_version.parse(flask_version) >= package_version.parse("3.1.0"):
+    # Flask 3.1+ introduced changes to request context handling
+    def _request_ctx_ref() -> weakref.ReferenceType:
+        try:
+            return weakref.ref(flask.globals.request_ctx._get_current_object())
+        except (RuntimeError, AttributeError):
+            # Handle cases where request context is not available or has changed
+            return weakref.ref(None)
+
+elif package_version.parse(flask_version) >= package_version.parse("2.2.0"):
 
     def _request_ctx_ref() -> weakref.ReferenceType:
         return weakref.ref(flask.globals.request_ctx._get_current_object())
@@ -333,6 +342,91 @@ def get_default_span_name():
     return span_name
 
 
+def _should_trace_request(excluded_urls) -> bool:
+    """Check if request should be traced based on excluded URLs."""
+    return bool(
+        flask.request
+        and (
+            excluded_urls is None
+            or not excluded_urls.url_disabled(flask.request.url)
+        )
+    )
+
+
+def _handle_response_headers(
+    status,
+    response_headers,
+    attributes,
+    span,
+    response_hook,
+    sem_conv_opt_in_mode,
+):
+    """Handle response headers and span attributes."""
+    propagator = get_global_response_propagator()
+    if propagator:
+        propagator.inject(
+            response_headers,
+            setter=otel_wsgi.default_response_propagation_setter,
+        )
+
+    if span:
+        otel_wsgi.add_response_attributes(
+            span,
+            status,
+            response_headers,
+            attributes,
+            sem_conv_opt_in_mode,
+        )
+        if span.is_recording() and span.kind == trace.SpanKind.SERVER:
+            custom_attributes = (
+                otel_wsgi.collect_custom_response_headers_attributes(
+                    response_headers
+                )
+            )
+            if len(custom_attributes) > 0:
+                span.set_attributes(custom_attributes)
+    else:
+        _logger.warning(
+            "Flask environ's OpenTelemetry span "
+            "missing at _start_response(%s)",
+            status,
+        )
+    if response_hook is not None:
+        response_hook(span, status, response_headers)
+
+
+def _record_metrics(
+    duration_s,
+    start_time,
+    request_route,
+    attributes,
+    duration_histogram_old,
+    duration_histogram_new,
+):
+    """Record duration metrics."""
+    if duration_histogram_old:
+        duration_attrs_old = otel_wsgi._parse_duration_attrs(
+            attributes, _StabilityMode.DEFAULT
+        )
+
+        if request_route:
+            # http.target to be included in old semantic conventions
+            duration_attrs_old[HTTP_TARGET] = str(request_route)
+
+        duration_histogram_old.record(
+            max(round(duration_s * 1000), 0), duration_attrs_old
+        )
+    if duration_histogram_new:
+        duration_attrs_new = otel_wsgi._parse_duration_attrs(
+            attributes, _StabilityMode.HTTP
+        )
+
+        if request_route:
+            duration_attrs_new[HTTP_ROUTE] = str(request_route)
+
+        duration_histogram_new.record(max(duration_s, 0), duration_attrs_new)
+
+
 def _rewrapped_app(
     wsgi_app,
     active_requests_counter,
@@ -361,89 +455,98 @@ def _wrapped_app(wrapped_app_environ, start_response):
 
         active_requests_counter.add(1, active_requests_count_attrs)
         request_route = None
-
         should_trace = True
 
         def _start_response(status, response_headers, *args, **kwargs):
-            nonlocal should_trace
-            should_trace = _should_trace(excluded_urls)
+            nonlocal should_trace, request_route
+            should_trace = _should_trace_request(excluded_urls)
+
             if should_trace:
-                nonlocal request_route
                 request_route = flask.request.url_rule
-
                 span = flask.request.environ.get(_ENVIRON_SPAN_KEY)
+                _handle_response_headers(
+                    status,
+                    response_headers,
+                    attributes,
+                    span,
+                    response_hook,
+                    sem_conv_opt_in_mode,
+                )
+            return start_response(status, response_headers, *args, **kwargs)
 
-                propagator = get_global_response_propagator()
-                if propagator:
-                    propagator.inject(
-                        response_headers,
-                        setter=otel_wsgi.default_response_propagation_setter,
-                    )
+        try:
+            result = wsgi_app(wrapped_app_environ, _start_response)
 
-                if span:
-                    otel_wsgi.add_response_attributes(
-                        span,
-                        status,
-                        response_headers,
-                        attributes,
-                        sem_conv_opt_in_mode,
-                    )
-                    if (
-                        span.is_recording()
-                        and span.kind == trace.SpanKind.SERVER
-                    ):
-                        custom_attributes = otel_wsgi.collect_custom_response_headers_attributes(
-                            response_headers
-                        )
-                        if len(custom_attributes) > 0:
-                            span.set_attributes(custom_attributes)
-                else:
-                    _logger.warning(
-                        "Flask environ's OpenTelemetry span "
-                        "missing at _start_response(%s)",
-                        status,
-                    )
-                if response_hook is not None:
-                    response_hook(span, status, response_headers)
-            return start_response(status, response_headers, *args, **kwargs)
+            # Handle streaming responses by ensuring proper cleanup
+            is_streaming = (
+                hasattr(result, "__iter__")
+                and not isinstance(result, (bytes, str))
+                and hasattr(result, "__next__")
+            )
 
-        result = wsgi_app(wrapped_app_environ, _start_response)
-        if should_trace:
-            duration_s = default_timer() - start
-            if duration_histogram_old:
-                duration_attrs_old = otel_wsgi._parse_duration_attrs(
-                    attributes, _StabilityMode.DEFAULT
+            if is_streaming:
+                # For streaming responses, defer cleanup until the response is consumed
+                # We'll use a weakref callback or rely on the teardown handler
+                pass
+            else:
+                # Non-streaming response, cleanup immediately
+                _cleanup_context_safely(wrapped_app_environ)
+
+            if should_trace:
+                duration_s = default_timer() - start
+                _record_metrics(
+                    duration_s,
+                    start,
+                    request_route,
+                    attributes,
+                    duration_histogram_old,
+                    duration_histogram_new,
                 )
+        except Exception:
+            # Ensure cleanup on exception
+            _cleanup_context_safely(wrapped_app_environ)
+            raise
+        finally:
+            active_requests_counter.add(-1, active_requests_count_attrs)
 
-                if request_route:
-                    # http.target to be included in old semantic conventions
-                    duration_attrs_old[HTTP_TARGET] = str(request_route)
+        return result
 
-                duration_histogram_old.record(
-                    max(round(duration_s * 1000), 0), duration_attrs_old
-                )
-            if duration_histogram_new:
-                duration_attrs_new = otel_wsgi._parse_duration_attrs(
-                    attributes, _StabilityMode.HTTP
-                )
+    def _cleanup_context_safely(wrapped_app_environ):
+        """Clean up context and tokens safely"""
+        try:
+            # Clean up activation and token to prevent context leaks
+            activation = wrapped_app_environ.get(_ENVIRON_ACTIVATION_KEY)
+            token = wrapped_app_environ.get(_ENVIRON_TOKEN)
+
+            if activation and hasattr(activation, "__exit__"):
+                try:
+                    activation.__exit__(None, None, None)
+                except (RuntimeError, AttributeError):
+                    _logger.debug(
+                        "Failed to exit activation during context cleanup",
+                        exc_info=True,
+                    )
 
-                if request_route:
-                    duration_attrs_new[HTTP_ROUTE] = str(request_route)
+            if token:
+                try:
+                    context.detach(token)
+                except (RuntimeError, AttributeError):
+                    _logger.debug(
+                        "Failed to detach token during context cleanup",
+                        exc_info=True,
+                    )
 
-                duration_histogram_new.record(
-                    max(duration_s, 0), duration_attrs_new
-                )
-        active_requests_counter.add(-1, active_requests_count_attrs)
-        return result
+            # Clean up environment keys
+            for key in [
+                _ENVIRON_ACTIVATION_KEY,
+                _ENVIRON_TOKEN,
+                _ENVIRON_SPAN_KEY,
+                _ENVIRON_REQCTX_REF_KEY,
+            ]:
+                wrapped_app_environ.pop(key, None)
 
-    def _should_trace(excluded_urls) -> bool:
-        return bool(
-            flask.request
-            and (
-                excluded_urls is None
-                or not excluded_urls.url_disabled(flask.request.url)
-            )
-        )
+        except (RuntimeError, AttributeError, KeyError):
+            _logger.debug("Exception during context cleanup", exc_info=True)
 
     return _wrapped_app
 
@@ -537,12 +640,26 @@ def _teardown_request(exc):
             return
 
         activation = flask.request.environ.get(_ENVIRON_ACTIVATION_KEY)
+        token = flask.request.environ.get(_ENVIRON_TOKEN)
+
+        # Check if this is a response that has already been cleaned up
+        if not activation and not token:
+            # Already cleaned up by streaming response handler
+            return
 
         original_reqctx_ref = flask.request.environ.get(
             _ENVIRON_REQCTX_REF_KEY
         )
-        current_reqctx_ref = _request_ctx_ref()
-        if not activation or original_reqctx_ref != current_reqctx_ref:
+
+        try:
+            current_reqctx_ref = _request_ctx_ref()
+        except (RuntimeError, AttributeError):
+            # Flask 3.1+ might raise exceptions when context is not available
+            current_reqctx_ref = None
+
+        if not activation or (
+            original_reqctx_ref and original_reqctx_ref != current_reqctx_ref
+        ):
             # This request didn't start a span, maybe because it was created in
             # a way that doesn't run `before_request`, like when it is created
             # with `app.test_request_context`.
@@ -554,15 +671,32 @@ def _teardown_request(exc):
             # like any decorated with `flask.copy_current_request_context`.
 
             return
-        if exc is None:
-            activation.__exit__(None, None, None)
-        else:
-            activation.__exit__(
-                type(exc), exc, getattr(exc, "__traceback__", None)
+
+        try:
+            if exc is None:
+                activation.__exit__(None, None, None)
+            else:
+                activation.__exit__(
+                    type(exc), exc, getattr(exc, "__traceback__", None)
+                )
+        except (RuntimeError, AttributeError) as teardown_exc:
+            _logger.debug(
+                "Failed to exit activation in teardown", exc_info=teardown_exc
+            )
+
+        try:
+            if token:
+                context.detach(token)
+        except (RuntimeError, AttributeError) as detach_exc:
+            _logger.debug(
+                "Failed to detach context in teardown", exc_info=detach_exc
             )
 
-        if flask.request.environ.get(_ENVIRON_TOKEN, None):
-            context.detach(flask.request.environ.get(_ENVIRON_TOKEN))
+        # Clean up environment keys to prevent memory leaks
+        flask.request.environ.pop(_ENVIRON_ACTIVATION_KEY, None)
+        flask.request.environ.pop(_ENVIRON_TOKEN, None)
+        flask.request.environ.pop(_ENVIRON_SPAN_KEY, None)
+        flask.request.environ.pop(_ENVIRON_REQCTX_REF_KEY, None)
 
     return _teardown_request