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

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