Decorator vs ASGI spect middleware execution issue with FastAPI

README.md

Decorator vs ASGI Spec Middleware with FastAPI and ddtrace Correlation

This illustrates something unexpected about middleware with FastAPI and ddtrace. Below is are two examples of the log output from the minimal application provided here. The first being the log output when the decorator style middleware is added. The second when the decorator style middleware is removed, notice the trace and span ID values are 0. The question is, what is going such that the trace/span ID is only available in the ASGI spec when the decorator style is also applied.

app.py

import logging
import structlog
import sys
import time
from typing import Dict, List, TypedDict
from asgiref.typing import ASGISendEvent, HTTPScope
from ddtrace import patch_all, tracer
from fastapi import FastAPI
from starlette.requests import Request
from starlette.types import ASGIApp, Receive, Send
from structlog.contextvars import merge_contextvars

patch_all()

def datadog_tracer_injection(logger, log_method, event_dict):
    # get correlation ids from current tracer context
    ctx = tracer.get_log_correlation_context()
    
    # add ids to structlog event dictionary
    event_dict['dd.trace_id'] = ctx["trace_id"]
    event_dict['dd.span_id'] = ctx["span_id"]

    # add the env, service, and version configured for the tracer
    event_dict['dd.env'] = ctx["env"]
    event_dict['dd.service'] = ctx["service"]
    event_dict['dd.version'] = ctx["version"]

    return event_dict

def rename_event_key(logger, method_name, event_dict):
    """Renames the `event` key to `message`

    This helper function renames the `event` key in structured logging
    entries to `message` key which conforms to Datadog's default
    attribute for log message text.
    """
    event_dict["message"] = event_dict.pop("event")
    return event_dict

def configure_logging() -> None:
    shared_processors = [
        merge_contextvars,
        structlog.stdlib.add_logger_name,
        structlog.stdlib.add_log_level,
        structlog.stdlib.PositionalArgumentsFormatter(),
        structlog.processors.TimeStamper(fmt="iso"),
        structlog.processors.StackInfoRenderer(),
        structlog.processors.format_exc_info,
        datadog_tracer_injection,
        rename_event_key
    ]
    
    renderer = structlog.processors.JSONRenderer()
        
    formatter = structlog.stdlib.ProcessorFormatter(
        processor=renderer, foreign_pre_chain=shared_processors
    )

    structlog_processors = [*shared_processors]
    structlog_processors.insert(0, structlog.stdlib.filter_by_level)
    structlog_processors.append(structlog.stdlib.ProcessorFormatter.wrap_for_formatter)

    structlog.configure(
        processors=structlog_processors,
        context_class=dict,
        logger_factory=structlog.stdlib.LoggerFactory(),
        wrapper_class=structlog.stdlib.BoundLogger,
    )

    handler = logging.StreamHandler(sys.stdout)
    handler.setFormatter(formatter)

    root_logger = logging.getLogger()
    root_logger.setLevel(logging.INFO)
    root_logger.handlers = []
    root_logger.addHandler(handler)

configure_logging()


logger = structlog.get_logger("app.request")

class AccessInfo(TypedDict, total=False):
    response: ASGISendEvent
    start_time: float
    end_time: float

class RequestLoggerMiddleware:
    def __init__(self, app: ASGIApp):
        self.app = app

    async def __call__(self, scope: HTTPScope, receive: Receive, send: Send):
        if scope["type"] != "http":
            return await self.app(scope, receive, send)

        info = AccessInfo(response={})

        async def inner_send(message: ASGISendEvent) -> None:
            if message["type"] == "http.response.start":
                info["response"] = message
            await send(message)

        try:
            info["start_time"] = time.monotonic_ns()
            await self.app(scope, receive, inner_send)
        except BaseException as exc:
            info["response"]["status"] = 500
            raise exc
        finally:
            info["end_time"] = time.monotonic_ns()
            logger.info("ASGI spec style")

app = FastAPI()

app.add_middleware(RequestLoggerMiddleware)

# If the following decorator style middleware is commented out the ddtrace trace/span ID's are NOT available to the ASGI style middleware
@app.middleware("logger")
async def add_process_time_header(request: Request, call_next):
    response = await call_next(request)
    logger.info("Decorator style")
    return response


@app.get("/")
async def index():
    logger.info("Index")
    return "hello"

uvicorn-with-decorator-middleware.log

$ poetry run uvicorn trace_test.app:app
INFO:     Started server process [24282]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
{"logger": "app.request", "level": "info", "timestamp": "2023-06-05T20:47:11.679314Z", "dd.trace_id": "5334270870132426139", "dd.span_id": "1449916201366737805", "dd.env": "", "dd.service": "fastapi", "dd.version": "", "message": "Index"}
{"logger": "app.request", "level": "info", "timestamp": "2023-06-05T20:47:11.679635Z", "dd.trace_id": "5334270870132426139", "dd.span_id": "1449916201366737805", "dd.env": "", "dd.service": "fastapi", "dd.version": "", "message": "Decorator style"}
INFO:     127.0.0.1:57712 - "GET / HTTP/1.1" 200 OK
{"logger": "app.request", "level": "info", "timestamp": "2023-06-05T20:47:11.680176Z", "dd.trace_id": "5334270870132426139", "dd.span_id": "1449916201366737805", "dd.env": "", "dd.service": "fastapi", "dd.version": "", "message": "ASGI spec style"}

uvicorn-without-decorator-middleware.log

$ poetry run uvicorn trace_test.app:app
INFO:     Started server process [24387]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
{"logger": "app.request", "level": "info", "timestamp": "2023-06-05T20:47:44.325602Z", "dd.trace_id": "7319887886799135321", "dd.span_id": "620939633164652408", "dd.env": "", "dd.service": "fastapi", "dd.version": "", "message": "Index"}
INFO:     127.0.0.1:57812 - "GET / HTTP/1.1" 200 OK
{"logger": "app.request", "level": "info", "timestamp": "2023-06-05T20:47:44.346619Z", "dd.trace_id": "0", "dd.span_id": "0", "dd.env": "", "dd.service": "", "dd.version": "", "message": "ASGI spec style"}