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())
@@ -361,12 +370,12 @@ 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)
+
             if should_trace:
                 nonlocal request_route
                 request_route = flask.request.url_rule
@@ -407,35 +416,84 @@ def _start_response(status, response_headers, *args, **kwargs):
                     response_hook(span, status, response_headers)
             return start_response(status, response_headers, *args, **kwargs)
 
-        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
-                )
+        try:
+            result = wsgi_app(wrapped_app_environ, _start_response)
 
-                if request_route:
-                    # http.target to be included in old semantic conventions
-                    duration_attrs_old[HTTP_TARGET] = str(request_route)
+            # Handle streaming responses by ensuring proper cleanup
+            is_streaming = (
+                hasattr(result, '__iter__') and
+                not isinstance(result, (bytes, str)) and
+                hasattr(result, '__next__')
+            )
 
-                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 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 request_route:
-                    duration_attrs_new[HTTP_ROUTE] = str(request_route)
+            if should_trace:
+                duration_s = default_timer() - start
+                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
+                    )
+        except Exception:
+            # Ensure cleanup on exception
+            _cleanup_context_safely(wrapped_app_environ)
+            raise
+        finally:
+            active_requests_counter.add(-1, active_requests_count_attrs)
 
-                duration_histogram_new.record(
-                    max(duration_s, 0), duration_attrs_new
-                )
-        active_requests_counter.add(-1, active_requests_count_attrs)
         return result
 
+    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 Exception:
+                    _logger.debug("Failed to exit activation during context cleanup", exc_info=True)
+
+            if token:
+                try:
+                    context.detach(token)
+                except Exception:
+                    _logger.debug("Failed to detach token during context cleanup", exc_info=True)
+
+            # 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)
+
+        except Exception:
+            _logger.debug("Exception during context cleanup", exc_info=True)
+
     def _should_trace(excluded_urls) -> bool:
         return bool(
             flask.request
@@ -537,12 +595,24 @@ 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 +624,28 @@ 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)
-            )
 
-        if flask.request.environ.get(_ENVIRON_TOKEN, None):
-            context.detach(flask.request.environ.get(_ENVIRON_TOKEN))
+        try:
+            if exc is None:
+                activation.__exit__(None, None, None)
+            else:
+                activation.__exit__(
+                    type(exc), exc, getattr(exc, "__traceback__", None)
+                )
+        except Exception as e:
+            _logger.debug("Failed to exit activation in teardown", exc_info=e)
+
+        try:
+            if token:
+                context.detach(token)
+        except Exception as e:
+            _logger.debug("Failed to detach context in teardown", exc_info=e)
+
+        # 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