fix(tracing): add flag for aiohttp that disables potential for memory leak (#12081)

This PR adds the environment variable
``DD_AIOHTTP_CLIENT_DISABLE_STREAM_TIMING_FOR_MEM_LEAK`` to address a
potential memory leak in the aiohttp integration. When set to true, this
flag may cause streamed response span timing to be inaccurate. The flag
defaults to false.

When this was merged #7551 we
essentially added support for getting the correct timing of streamed
responses, however it also caused us to sometimes never close spans,
which lead to a memory leak. When set to true, this flag essentially
reverts that behavior.

- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

- [x] Reviewer has checked that all the criteria below are met
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

(cherry picked from commit 17c60dd)

Author: ZStriker19

ddtrace/contrib/aiohttp/__init__.py

@@ -36,6 +36,13 @@
 
    Default: ``False``
 
+.. py:data:: ddtrace.config.aiohttp['disable_stream_timing_for_mem_leak']
+
+   Whether or not to to address a potential memory leak in the aiohttp integration.
+   When set to ``True``, this flag may cause streamed response span timing to be inaccurate.
+
+   Default: ``False``
+
 
 Server
 ******

ddtrace/contrib/internal/aiohttp/middlewares.py

@@ -63,19 +63,39 @@ async def attach_context(request):
         if (config._analytics_enabled and analytics_enabled is not False) or analytics_enabled is True:
             request_span.set_tag(_ANALYTICS_SAMPLE_RATE_KEY, app[CONFIG_KEY].get("analytics_sample_rate", True))
 
-        # attach the context and the root span to the request; the Context
-        # may be freely used by the application code
-        request[REQUEST_CONTEXT_KEY] = request_span.context
-        request[REQUEST_SPAN_KEY] = request_span
-        request[REQUEST_CONFIG_KEY] = app[CONFIG_KEY]
-        try:
-            response = await handler(request)
-            if isinstance(response, web.StreamResponse):
-                request.task.add_done_callback(lambda _: finish_request_span(request, response))
-            return response
-        except Exception:
-            request_span.set_traceback()
-            raise
+        with core.context_with_data(
+            "aiohttp.request",
+            span_name=schematize_url_operation("aiohttp.request", protocol="http", direction=SpanDirection.INBOUND),
+            span_type=SpanTypes.WEB,
+            service=service,
+            tags={},
+            tracer=tracer,
+            distributed_headers=request.headers,
+            distributed_headers_config=config.aiohttp,
+            distributed_headers_config_override=app[CONFIG_KEY]["distributed_tracing_enabled"],
+            headers_case_sensitive=True,
+            analytics_enabled=analytics_enabled,
+            analytics_sample_rate=app[CONFIG_KEY].get("analytics_sample_rate", True),
+        ) as ctx:
+            req_span = ctx.span
+
+            ctx.set_item("req_span", req_span)
+            core.dispatch("web.request.start", (ctx, config.aiohttp))
+
+            # attach the context and the root span to the request; the Context
+            # may be freely used by the application code
+            request[REQUEST_CONTEXT_KEY] = req_span.context
+            request[REQUEST_SPAN_KEY] = req_span
+            request[REQUEST_CONFIG_KEY] = app[CONFIG_KEY]
+            try:
+                response = await handler(request)
+                if not config.aiohttp["disable_stream_timing_for_mem_leak"]:
+                    if isinstance(response, web.StreamResponse):
+                        request.task.add_done_callback(lambda _: finish_request_span(request, response))
+                return response
+            except Exception:
+                req_span.set_traceback()
+                raise
 
     return attach_context
 
@@ -142,9 +162,13 @@ async def on_prepare(request, response):
     the trace middleware execution.
     """
     # NB isinstance is not appropriate here because StreamResponse is a parent of the other
-    # aiohttp response types
-    if type(response) is web.StreamResponse and not response.task.done():
-        return
+    # aiohttp response types. However in some cases this can also lead to missing the closing of
+    # spans, leading to a memory leak, which is why we have this flag.
+    # todo: this is a temporary fix for a memory leak in aiohttp. We should find a way to
+    # consistently close spans with the correct timing.
+    if not config.aiohttp["disable_stream_timing_for_mem_leak"]:
+        if type(response) is web.StreamResponse and not response.task.done():
+            return
     finish_request_span(request, response)
 
 

ddtrace/contrib/internal/aiohttp/patch.py

@@ -22,6 +22,7 @@
 from ddtrace.internal.utils import get_argument_value
 from ddtrace.internal.utils.formats import asbool
 from ddtrace.propagation.http import HTTPPropagator
+from ddtrace.settings._core import get_config as _get_config
 from ddtrace.trace import Pin
 
 
@@ -31,7 +32,12 @@
 # Server config
 config._add(
     "aiohttp",
-    dict(distributed_tracing=True),
+    dict(
+        distributed_tracing=True,
+        disable_stream_timing_for_mem_leak=asbool(
+            _get_config("DD_AIOHTTP_CLIENT_DISABLE_STREAM_TIMING_FOR_MEM_LEAK", default=False)
+        ),
+    ),
 )
 
 config._add(

releasenotes/notes/add_aiohttp_memory_leak_flag-66005f987dbfbd47.yaml

@@ -0,0 +1,5 @@
+---
+
+fixes:
+  - |
+    aiohttp: Adds the environment variable ``DD_AIOHTTP_CLIENT_DISABLE_STREAM_TIMING_FOR_MEM_LEAK`` to address a potential memory leak in the aiohttp integration. When set to true, this flag may cause streamed response span timing to be inaccurate. The flag defaults to false.

tests/contrib/aiohttp/test_request.py

@@ -31,6 +31,33 @@ async def test_full_request(patched_app_tracer, aiohttp_client, loop):
     assert "GET /" == request_span.resource
 
 
+async def test_full_request_w_mem_leak_prevention_flag(patched_app_tracer, aiohttp_client, loop):
+    config.aiohttp.disable_stream_timing_for_mem_leak = True
+    try:
+        app, tracer = patched_app_tracer
+        client = await aiohttp_client(app)
+        # it should create a root span when there is a handler hit
+        # with the proper tags
+        request = await client.request("GET", "/")
+        assert 200 == request.status
+        await request.text()
+        # the trace is created
+        traces = tracer.pop_traces()
+        assert 1 == len(traces)
+        assert 1 == len(traces[0])
+        request_span = traces[0][0]
+        assert_is_measured(request_span)
+
+        # request
+        assert "aiohttp-web" == request_span.service
+        assert "aiohttp.request" == request_span.name
+        assert "GET /" == request_span.resource
+    except Exception:
+        raise
+    finally:
+        config.aiohttp.disable_stream_timing_for_mem_leak = False
+
+
 async def test_stream_request(patched_app_tracer, aiohttp_client, loop):
     app, tracer = patched_app_tracer
     async with await aiohttp_client(app) as client: