Repository: DataDog/dd-trace-py

Issue:
Starlette/Fastapi: endpoint duration includes the duration of background tasks
### Which version of dd-trace-py are you using?

ddtrace==0.55.4

### Which version of pip are you using?

21.2.4


### Which version of the libraries are you using?

fastapi==0.68.2
starlette==0.14.2

### How can we reproduce your problem?

this would be a minimal proof of concept `app.py`, running through `ddtrace-run uvicorn app:app`

```
import asyncio

from ddtrace import tracer
from fastapi import FastAPI, BackgroundTasks

app = FastAPI()


async def some_background_task():
    with tracer.start_span("some_background_task", activate=True):
        tracer.context_provider.activate(None)
        await asyncio.sleep(10)


@app.get("/")
async def main(background_tasks: BackgroundTasks):
    background_tasks.add_task(some_background_task)
    return "ok"

```

### What is the result that you get?

The duration of `/` is reported to be 10s, while the browser immediately receives the response.
`some_background_task` is also reported with a duration of 10s.

### What is the result that you expected?

I would expect that the reported endpoint duration matches the time it took to get the response, and that the background task is reported separately. Please don't mind that `tracer.context_provider.activate(None)` might be redundant here, adding it here to show what I have tried.

FastAPI's `add_task` actually comes from starlette https://www.starlette.io/background/

I can understand why the endpoint duration includes the background task, this is the definition of starlette's `Response.__call__`:

https://github.com/encode/starlette/blob/ada99beee530e7b841ce320bc6e66f6dbd9ad781/starlette/responses.py#L159
```
    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        await send(
            {
                "type": "http.response.start",
                "status": self.status_code,
                "headers": self.raw_headers,
            }
        )
        await send({"type": "http.response.body", "body": self.body})

        if self.background is not None:
            await self.background()
```

The response header and body is sent, but the function itself is not finished until all background tasks have been processed.

I believe that this is not what users of ddtrace would expect: the background tasks are used to return a response early without waiting for background operations to finish ; the reported endpoint duration should correspond to when the body was sent



Discussion:
Hello, any updates on this problem?

Resolving patch:
diff --git a/ddtrace/contrib/asgi/middleware.py b/ddtrace/contrib/asgi/middleware.py
--- a/ddtrace/contrib/asgi/middleware.py
+++ b/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()


Test patch:
diff --git a/tests/contrib/asgi/test_asgi.py b/tests/contrib/asgi/test_asgi.py
--- a/tests/contrib/asgi/test_asgi.py
+++ b/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
diff --git a/tests/contrib/fastapi/app.py b/tests/contrib/fastapi/app.py
--- a/tests/contrib/fastapi/app.py
+++ b/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}")
diff --git a/tests/contrib/fastapi/test_fastapi.py b/tests/contrib/fastapi/test_fastapi.py
--- a/tests/contrib/fastapi/test_fastapi.py
+++ b/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
diff --git a/tests/contrib/starlette/app.py b/tests/contrib/starlette/app.py
--- a/tests/contrib/starlette/app.py
+++ b/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])
diff --git a/tests/contrib/starlette/test_starlette.py b/tests/contrib/starlette/test_starlette.py
--- a/tests/contrib/starlette/test_starlette.py
+++ b/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