fix(asgi, starlette, fastapi): exclude background task durations from web requests (#3799)

Following up on #3023, currently libraries traced via the `asgi` integration, such as`fastapi.request` and `starlette.request` spans include the duration of background tasks. This PR updates the asgi middleware code so that spans end earlier by excluding background tasks to better represent the web requests response times.

**Current Behavior**
$$\begin{align}duration = {http response time} + {background tasks time}\end{align}$$

**What it looks like with the PR**
$$\begin{align}duration = {http response time}\end{align}$$

**Example:** An application return a "hello world!" to the user in $2 milliseconds$, but the events that happen in the background lasts another $8 seconds$.

With this PR, the span will report a duration of $2 milliseconds$, instead of $2 milliseconds + 8 seconds$.

Fixes #3023

Also fixes this failing test: https://app.circleci.com/pipelines/github/DataDog/dd-trace-py/18150/workflows/24caa985-5805-4bff-a4eb-5b143f24a2fb/jobs/1233079

## Checklist
- [ ] Library documentation is updated.
- [ ] [Corp site](https://github.com/DataDog/documentation/) documentation is updated (link to the PR).

## Reviewer Checklist
- [ ] Title is accurate.
- [ ] Description motivates each change.
- [ ] No unnecessary changes were introduced in this PR.
- [ ] PR cannot be broken up into smaller PRs.
- [ ] Avoid breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes unless absolutely necessary.
- [ ] Tests provided or description of manual testing performed is included in the code or PR.
- [ ] Release note has been added for fixes and features, or else `changelog/no-changelog` label added.
- [ ] All relevant GitHub issues are correctly linked.
- [ ] Backports are identified and tagged with Mergifyio.
- [ ] Add to milestone.

(cherry picked from commit 5b97489)

Author: wantsui

ddtrace/contrib/asgi/middleware.py

@@ -169,7 +169,13 @@ async def wrapped_send(message):
                 span, self.integration_config, status_code=status_code, response_headers=response_headers
             )
 
-            return await send(message)
+            try:
+                return await send(message)
+            finally:
+                # Per asgi spec, "more_body" is used if there is still data to send
+                # Close the span if "http.response.body" has no more data left to send in the response.
+                if message.get("type") == "http.response.body" and not message.get("more_body", False):
+                    span.finish()
 
         try:
             return await self.app(scope, receive, wrapped_send)
@@ -183,4 +189,5 @@ async def wrapped_send(message):
                 del scope["datadog"]["request_span"]
             except KeyError:
                 pass
+
             span.finish()

releasenotes/notes/fix-asgi-durations-include-background-tasks-4add6b50712608b2.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    asgi, starlette, fastapi: Exclude background tasks duration from web request spans.

tests/contrib/asgi/test_asgi.py

@@ -94,6 +94,44 @@ def double_callable_app(scope):
     return partial(basic_app, scope)
 
 
+async def tasks_app_without_more_body(scope, receive, send):
+    """
+    An app that does something in the background after the response is sent without having more data to send.
+    "more_body" with a true value is used in the asgi spec to indicate that there is more data to send.
+    """
+    assert scope["type"] == "http"
+    message = await receive()
+    if message.get("type") == "http.request":
+        await send({"type": "http.response.start", "status": 200, "headers": [[b"Content-Type", b"text/plain"]]})
+        await send({"type": "http.response.body", "body": b"*"})
+        await asyncio.sleep(1)
+
+
+async def tasks_app_with_more_body(scope, receive, send):
+    """
+    An app that does something in the background but has a more_body response that starts off as true,
+    but then turns to false.
+    "more_body" with a true value is used in the asgi spec to indicate that there is more data to send.
+    """
+    assert scope["type"] == "http"
+    message = await receive()
+    request_span = scope["datadog"]["request_spans"][0]
+    if message.get("type") == "http.request":
+
+        # assert that the request span hasn't finished at the start of a response
+        await send({"type": "http.response.start", "status": 200, "headers": [[b"Content-Type", b"text/plain"]]})
+        assert not request_span.finished
+
+        # assert that the request span hasn't finished while more_body is True
+        await send({"type": "http.response.body", "body": b"*", "more_body": True})
+        assert not request_span.finished
+
+        # assert that the span has finished after more_body is False
+        await send({"type": "http.response.body", "body": b"*", "more_body": False})
+        assert request_span.finished
+        await asyncio.sleep(1)
+
+
 def _check_span_tags(scope, span):
     assert span.get_tag("http.method") == scope["method"]
     server = scope.get("server")
@@ -445,3 +483,49 @@ async def test_app_no_middleware(scope, receive, send):
     async with httpx.AsyncClient(app=test_app_no_middleware) as client:
         response = await client.get("http://testserver/")
         assert response.status_code == 200
+
+
+@pytest.mark.asyncio
+async def test_tasks_asgi_without_more_body(scope, tracer, test_spans):
+    """
+    When an application doesn't have more_body calls and does background tasks,
+    the asgi span only captures the time it took for a user to get a response, not the time
+    it took for other tasks in the background.
+    """
+    app = TraceMiddleware(tasks_app_without_more_body, tracer=tracer)
+    async with httpx.AsyncClient(app=app) as client:
+        response = await client.get("http://testserver/")
+        assert response.status_code == 200
+
+    spans = test_spans.pop_traces()
+    assert len(spans) == 1
+    assert len(spans[0]) == 1
+    request_span = spans[0][0]
+    assert request_span.name == "asgi.request"
+    assert request_span.span_type == "web"
+    # typical duration without background task should be in less than 10 ms
+    # duration with background task will take approximately 1.1s
+    assert request_span.duration < 1
+
+
+@pytest.mark.asyncio
+async def test_tasks_asgi_with_more_body(scope, tracer, test_spans):
+    """
+    When an application does have more_body calls and does background tasks,
+    the asgi span only captures the time it took for a user to get a response, not the time
+    it took for other tasks in the background.
+    """
+    app = TraceMiddleware(tasks_app_with_more_body, tracer=tracer)
+    async with httpx.AsyncClient(app=app) as client:
+        response = await client.get("http://testserver/")
+        assert response.status_code == 200
+
+    spans = test_spans.pop_traces()
+    assert len(spans) == 1
+    assert len(spans[0]) == 1
+    request_span = spans[0][0]
+    assert request_span.name == "asgi.request"
+    assert request_span.span_type == "web"
+    # typical duration without background task should be in less than 10 ms
+    # duration with background task will take approximately 1.1s
+    assert request_span.duration < 1

tests/contrib/fastapi/app.py

@@ -1,7 +1,9 @@
+import asyncio
 from tempfile import NamedTemporaryFile
 import time
 from typing import Optional
 
+from fastapi import BackgroundTasks
 from fastapi import FastAPI
 from fastapi import HTTPException
 from fastapi import Header
@@ -102,6 +104,14 @@ async def file():
             fp.flush()
             return FileResponse(fp.name)
 
+    async def custom_task():
+        await asyncio.sleep(1)
+
+    @app.get("/asynctask")
+    async def asynctask(bg_tasks: BackgroundTasks):
+        bg_tasks.add_task(custom_task)
+        return "task added"
+
     subapp = FastAPI()
 
     @subapp.get("/hello/{name}")

tests/contrib/fastapi/test_fastapi.py

@@ -552,3 +552,20 @@ def test_table_query_snapshot(snapshot_client):
         "name": "Test Name",
         "description": "This request adds a new entry to the test db",
     }
+
+
+def test_background_task(client, tracer, test_spans):
+    """Tests if background tasks have been excluded from span duration"""
+    response = client.get("/asynctask")
+    assert response.status_code == 200
+    assert response.json() == "task added"
+    spans = test_spans.pop_traces()
+    assert len(spans) == 1
+    assert len(spans[0]) == 2
+    request_span, serialize_span = spans[0]
+
+    assert request_span.name == "fastapi.request"
+    assert request_span.resource == "GET /asynctask"
+    # typical duration without background task should be in less than 10 ms
+    # duration with background task will take approximately 1.1s
+    assert request_span.duration < 1

tests/contrib/starlette/app.py

@@ -1,10 +1,13 @@
+import asyncio
 from tempfile import NamedTemporaryFile
 import time
 
 import databases
 import sqlalchemy
 from starlette.applications import Starlette
+from starlette.background import BackgroundTask
 from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
 from starlette.responses import PlainTextResponse
 from starlette.responses import Response
 from starlette.responses import StreamingResponse
@@ -106,6 +109,14 @@ async def add_note(request):
         response = "Success"
         return PlainTextResponse(response)
 
+    async def custom_task():
+        await asyncio.sleep(1)
+
+    async def background_task(request):
+        """An example endpoint that just adds to background tasks"""
+        jsonmsg = {"result": "Background task added"}
+        return JSONResponse(jsonmsg, background=BackgroundTask(custom_task))
+
     routes = [
         Route("/", endpoint=homepage, name="homepage", methods=["GET"]),
         Route("/200", endpoint=success, name="200", methods=["GET"]),
@@ -119,6 +130,7 @@ async def add_note(request):
         Route("/notes", endpoint=list_notes, methods=["GET"]),
         Route("/notes", endpoint=add_note, methods=["POST"]),
         Mount("/sub-app", Starlette(routes=[Route("/hello/{name}", endpoint=success, name="200", methods=["GET"])])),
+        Route("/backgroundtask", endpoint=background_task, name="200", methods=["GET"]),
     ]
 
     app = Starlette(routes=routes, on_startup=[database.connect], on_shutdown=[database.disconnect])

tests/contrib/starlette/test_starlette.py

@@ -450,3 +450,17 @@ def test_incorrect_patching(run_python_code_in_subprocess):
     assert status == 0, err
     assert out == b"", err
     assert err == b"datadog context not present in ASGI request scope, trace middleware may be missing\n"
+
+
+def test_background_task(client, tracer, test_spans):
+    """Tests if background tasks have been excluded from span duration"""
+    r = client.get("/backgroundtask")
+    assert r.status_code == 200
+    assert r.text == '{"result":"Background task added"}'
+
+    request_span = next(test_spans.filter_spans(name="starlette.request"))
+    assert request_span.name == "starlette.request"
+    assert request_span.resource == "GET /backgroundtask"
+    # typical duration without background task should be in less than 10ms
+    # duration with background task will take approximately 1.1s
+    assert request_span.duration < 1