[Bug]FastAPI request: got different trace_ids at the error handler

Search before asking

[X] I had searched in the issues and found no similar issues.

Apache SkyWalking Component

Python Agent (apache/skywalking-python)

What happened

I want to return the skywalking trace id at the error handlers in FastAPI.

I print the trace id at the beginning of the request and in the error handler, and I got two different trace ids.(The first one is correct, which can be search in the skywalking UI)

For example, I got the log with trace id 8f954eae8f4011ee9a7682b792d8f7f3:

Exception: foo error: 8f954eae8f4011ee9a7682b792d8f7f3

And the response of the request with trace id 8f98f8108f4011ee9a7682b792d8f7f3:

"trace_id: 8f98f8108f4011ee9a7682b792d8f7f3"

What you expected to happen

I expect the two trace ids should be same. By the way, I get trace id from context:

def get_trace_id() -> str:
    """get skywalking trace_id"""
    try:
        context = get_context()
        trace_id = str(context.segment.related_traces[0])
        return trace_id
    except Exception as e:
        return f"get trace id error: {e}"

How to reproduce

Run the demo app service(sw-python run uvicorn app:app --reload) and call the demo api:

import logging

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

from skywalking.trace.context import get_context

logger = logging.getLogger(__name__)

def get_trace_id() -> str:
    """get skywalking trace_id"""
    try:
        context = get_context()
        trace_id = str(context.segment.related_traces[0])
        return trace_id
    except Exception as e:
        return f"get trace id error: {e}"

app = FastAPI()

@app.get("/")
def read_root():
    return {"Hello": "World"}

@app.get("/hello")
def hello():
    trace_id = get_trace_id()
    logger.info(f"request start: trace_id: {trace_id}")
    raise Exception(f"foo error: {trace_id}")

@app.exception_handler(Exception)
def handle_http_exception(request: Request, exc: Exception):
    logger.info(f"error handler: trace_id: {get_trace_id()}")
    return JSONResponse(
        content=f"trace_id: {get_trace_id()}",
    )

Anything else

I try to add a http middleware to FastAPI and use request.state.traceid to save the trace id, but it doesn't work, I still got different trace ids.

Finally, I find a way to solve it partly: at every entry point and get the state at the error handler. The method works, but it's verbose, which means that I have add all my entry point function with same code: a request parameter and the state save process. So I want to know is there any other method to do it elegantly.

@app.get("/hello")
def hello(request: Request):
    trace_id = get_trace_id()
    request.state.traceid = trace_id. # HERE!
    logger.info(f"request start: trace_id: {trace_id}")
    raise Exception(f"foo error: {trace_id}")

The environments(partly)

export SW_AGENT_PROTOCOL=http
export SW_AGENT_LOG_REPORTER_LEVEL=INFO
export SW_PLUGIN_PYMONGO_TRACE_PARAMETERS=True

Are you willing to submit a pull request to fix on your own?

[ ] Yes I am willing to submit a pull request on my own!

Code of Conduct

[X] I agree to follow this project's Code of Conduct

apache / skywalking