Use request_ctx to determine whether or not _teardown_request should end flask span

dev-requirements.txt

@@ -14,7 +14,7 @@ bleach==4.1.0 # transient dependency for readme-renderer
 grpcio-tools==1.29.0
 mypy-protobuf>=1.23
 protobuf~=3.13
-markupsafe==2.0.1
+markupsafe>=2.0.1
 codespell==2.1.0
 requests==2.28.1
 ruamel.yaml==0.17.21

instrumentation/opentelemetry-instrumentation-flask/pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
   "opentelemetry-instrumentation-wsgi == 0.39b0.dev",
   "opentelemetry-semantic-conventions == 0.39b0.dev",
   "opentelemetry-util-http == 0.39b0.dev",
+  "packaging >= 21.0",
 ]
 
 [project.optional-dependencies]
@@ -38,7 +39,7 @@ instruments = [
 ]
 test = [
   "opentelemetry-instrumentation-flask[instruments]",
-  "markupsafe==2.0.1",
+  "markupsafe>=2.0.1",
   "opentelemetry-test-utils == 0.39b0.dev",
 ]
 

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

@@ -239,10 +239,11 @@ def response_hook(span: Span, status: str, response_headers: List):
 ---
 """
 from logging import getLogger
-from threading import get_ident
+from packaging import version as package_version
 from time import time_ns
 from timeit import default_timer
 from typing import Collection
+import weakref
 
 import flask
 
@@ -265,11 +266,17 @@ def response_hook(span: Span, status: str, response_headers: List):
 _ENVIRON_STARTTIME_KEY = "opentelemetry-flask.starttime_key"
 _ENVIRON_SPAN_KEY = "opentelemetry-flask.span_key"
 _ENVIRON_ACTIVATION_KEY = "opentelemetry-flask.activation_key"
-_ENVIRON_THREAD_ID_KEY = "opentelemetry-flask.thread_id_key"
+_ENVIRON_REQCTX_REF_KEY = "opentelemetry-flask.reqctx_ref_key"
 _ENVIRON_TOKEN = "opentelemetry-flask.token"
 
 _excluded_urls_from_env = get_excluded_urls("FLASK")
 
+if 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())
+else:
+    def _request_ctx_ref() -> weakref.ReferenceType:
+        return weakref.ref(flask._request_ctx_stack.top)
 
 def get_default_span_name():
     try:
@@ -399,7 +406,7 @@ def _before_request():
         activation = trace.use_span(span, end_on_exit=True)
         activation.__enter__()  # pylint: disable=E1101
         flask_request_environ[_ENVIRON_ACTIVATION_KEY] = activation
-        flask_request_environ[_ENVIRON_THREAD_ID_KEY] = get_ident()
+        flask_request_environ[_ENVIRON_REQCTX_REF_KEY] = _request_ctx_ref()
         flask_request_environ[_ENVIRON_SPAN_KEY] = span
         flask_request_environ[_ENVIRON_TOKEN] = token
 
@@ -439,17 +446,20 @@ def _teardown_request(exc):
             return
 
         activation = flask.request.environ.get(_ENVIRON_ACTIVATION_KEY)
-        thread_id = flask.request.environ.get(_ENVIRON_THREAD_ID_KEY)
-        if not activation or thread_id != get_ident():
+
+        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:
             # 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`.
             #
-            # Similarly, check the thread_id against the current thread to ensure
-            # tear down only happens on the original thread. This situation can
-            # arise if the original thread handling the request spawn children
-            # threads and then uses something like copy_current_request_context
-            # to copy the request context.
+            # Similarly, check that the request_ctx that created the span
+            # matches the current request_ctx, and only tear down if they match.
+            # This situation can arise if the original request_ctx handling
+            # the request calls functions that push new request_ctx's,
+            # like any decorated with `flask.copy_current_request_context`.
+
             return
         if exc is None:
             activation.__exit__(None, None, None)

instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py

@@ -20,6 +20,7 @@
 from werkzeug.wrappers import Response
 
 from opentelemetry import context
+from opentelemetry import trace
 
 
 class InstrumentationTest:
@@ -37,6 +38,22 @@ def _sqlcommenter_endpoint():
         )
         return sqlcommenter_flask_values
 
+    @staticmethod
+    def _copy_context_endpoint():
+
+        @flask.copy_current_request_context
+        def _extract_header():
+            return flask.request.headers['x-req']
+
+        # Despite `_extract_header` copying the request context,
+        # calling it shouldn't detach the parent Flask span's contextvar
+        request_header = _extract_header()
+
+        return {
+            'span_name': trace.get_current_span().name,
+            'request_header': request_header
+        }
+
     @staticmethod
     def _multithreaded_endpoint(count):
         def do_random_stuff():
@@ -84,6 +101,7 @@ def excluded2_endpoint():
         self.app.route("/hello/<int:helloid>")(self._hello_endpoint)
         self.app.route("/sqlcommenter")(self._sqlcommenter_endpoint)
         self.app.route("/multithreaded")(self._multithreaded_endpoint)
+        self.app.route("/copy_context")(self._copy_context_endpoint)
         self.app.route("/excluded/<int:helloid>")(self._hello_endpoint)
         self.app.route("/excluded")(excluded_endpoint)
         self.app.route("/excluded2")(excluded2_endpoint)

instrumentation/opentelemetry-instrumentation-flask/tests/test_copy_context.py

@@ -0,0 +1,52 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import flask
+from werkzeug.test import Client
+from werkzeug.wrappers import Response
+
+from opentelemetry.instrumentation.flask import FlaskInstrumentor
+from opentelemetry.test.wsgitestutil import WsgiTestBase
+
+# pylint: disable=import-error
+from .base_test import InstrumentationTest
+
+
+class TestCopyContext(InstrumentationTest, WsgiTestBase):
+    def setUp(self):
+        super().setUp()
+        FlaskInstrumentor().instrument()
+        self.app = flask.Flask(__name__)
+        self._common_initialization()
+
+    def tearDown(self):
+        super().tearDown()
+        with self.disable_logging():
+            FlaskInstrumentor().uninstrument()
+
+    def test_copycontext(self):
+        """Test that instrumentation tear down does not blow up
+        when the request calls functions where the context has been
+        copied via `flask.copy_current_request_context`
+        """
+        self.app = flask.Flask(__name__)
+        self.app.route("/copy_context")(
+            self._copy_context_endpoint
+        )
+        client = Client(self.app, Response)
+        resp = client.get("/copy_context", headers={'x-req': 'a-header'})
+
+        self.assertEqual(200, resp.status_code)
+        self.assertEqual('/copy_context', resp.json['span_name'])
+        self.assertEqual('a-header', resp.json['request_header'])

tox.ini

@@ -80,8 +80,8 @@ envlist =
     pypy3-test-instrumentation-fastapi
 
     ; opentelemetry-instrumentation-flask
-    py3{7,8,9,10,11}-test-instrumentation-flask
-    pypy3-test-instrumentation-flask
+    py3{7,8,9,10,11}-test-instrumentation-flask{213,220}
+    pypy3-test-instrumentation-flask{213,220}
 
     ; opentelemetry-instrumentation-urllib
     py3{7,8,9,10,11}-test-instrumentation-urllib
@@ -254,6 +254,8 @@ deps =
   falcon1: falcon ==1.4.1
   falcon2: falcon >=2.0.0,<3.0.0
   falcon3: falcon >=3.0.0,<4.0.0
+  flask213: Flask ==2.1.3
+  flask220: Flask >=2.2.0
   grpc: pytest-asyncio
   sqlalchemy11: sqlalchemy>=1.1,<1.2
   sqlalchemy14: aiosqlite
@@ -359,8 +361,8 @@ commands_pre =
 
   grpc: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-grpc[test]
 
-  falcon{1,2,3},flask,django{1,2,3,4},pyramid,tornado,starlette,fastapi,aiohttp,asgi,requests,urllib,urllib3,wsgi: pip install {toxinidir}/util/opentelemetry-util-http[test]
-  wsgi,falcon{1,2,3},flask,django{1,2,3,4},pyramid: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-wsgi[test]
+  falcon{1,2,3},flask{213,220},django{1,2,3,4},pyramid,tornado,starlette,fastapi,aiohttp,asgi,requests,urllib,urllib3,wsgi: pip install {toxinidir}/util/opentelemetry-util-http[test]
+  wsgi,falcon{1,2,3},flask{213,220},django{1,2,3,4},pyramid: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-wsgi[test]
   asgi,django{3,4},starlette,fastapi: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-asgi[test]
 
   asyncpg: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-asyncpg[test]
@@ -374,7 +376,7 @@ commands_pre =
 
   falcon{1,2,3}: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-falcon[test]
 
-  flask: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-flask[test]
+  flask{213,220}: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-flask[test]
 
   urllib: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-urllib[test]