fix(tracing): add flag for aiohttp that disables potential for memory leak (#12081) [backport 2.21] (#12335)

Backports #12081 to 2.21. I
had to make some manual adjustments because there was a big change
between 2.x and 3.x.

From @ZStriker19 's PR description:
> 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.


(cherry picked from commit 17c60dd)

## Checklist
- [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))

## Reviewer Checklist
- [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)

---------

Co-authored-by: Zachary Groves <[email protected]>

Author: wantsui

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

@@ -70,8 +70,9 @@ async def attach_context(request):
         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))
+            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:
             request_span.set_traceback()
@@ -142,9 +143,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: