diff --git a/ddtrace/contrib/asgi/middleware.py b/ddtrace/contrib/asgi/middleware.py index 85234f282ba..8b97b04ea80 100644 --- 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() diff --git a/releasenotes/notes/fix-asgi-durations-include-background-tasks-4add6b50712608b2.yaml b/releasenotes/notes/fix-asgi-durations-include-background-tasks-4add6b50712608b2.yaml new file mode 100644 index 00000000000..345e9547ba8 --- /dev/null +++ b/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. diff --git a/tests/contrib/asgi/test_asgi.py b/tests/contrib/asgi/test_asgi.py index 1ed6153b5fa..5d7758586b6 100644 --- 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 index ddb1f77fb5f..24e7830c2ed 100644 --- 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 index bf013fbc8ea..acc1322e055 100644 --- 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 index b787e3914a1..40885ad55a4 100644 --- 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 index 9a53f130bd7..5ffb93f3a50 100644 --- 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