kaxil commented on code in PR #68039:
URL: https://github.com/apache/airflow/pull/68039#discussion_r3359357859
##########
airflow-core/src/airflow/api_fastapi/execution_api/app.py:
##########
@@ -253,9 +252,16 @@ async def _extract_w3c_trace_context(
token = otel_context.attach(ctx)
try:
yield
- finally:
- with contextlib.suppress(Exception):
- otel_context.detach(token)
+ except GeneratorExit:
+ # The dependency generator was force-closed (client disconnect /
request
+ # cancellation): the finalizer runs in a different asyncio Task — and
thus a
+ # different contextvars.Context — than attach did, so detaching the
token would
+ # raise "Token was created in a different Context" (which OTel logs at
ERROR
+ # before any suppression here could see it). The attached Context is
being
+ # discarded with the dying task, so detaching is unnecessary; skip it
and re-raise.
+ raise
+ else:
+ otel_context.detach(token)
Review Comment:
The `else` only detaches on normal completion, but there's a third unwind
path this misses. When a route handler raises (not a disconnect), FastAPI
finalizes this dependency by `athrow`-ing the exception into the generator at
the `yield` (via `AsyncExitStack` / `asynccontextmanager`, see
`fastapi/dependencies/utils.py` `_solve_generator`). That takes neither `except
GeneratorExit` nor `else`, so detach is skipped. Unlike the disconnect case,
`attach` here ran in the *same* task, so detach would have succeeded (the old
`finally` did it), and the W3C context now stays attached for the rest of the
request task: the `@exception_handler`, the 500 response, and any spans/logs
emitted during error handling inherit the stale upstream trace context. Could
you detach on same-task unwind paths too, e.g. add `except BaseException:
otel_context.detach(token); raise` alongside the `GeneratorExit` skip
(GeneratorExit is a `BaseException`, so order it first)?
##########
airflow-core/tests/unit/api_fastapi/execution_api/test_app.py:
##########
@@ -244,3 +247,34 @@ async def mock_require_auth(request: Request) -> TIToken:
):
with pytest.raises(RuntimeError, match="boom"):
test_client.get("/variables/k", headers={"Authorization":
"Bearer fake"})
+
+ @staticmethod
+ def _make_request() -> Request:
+ return Request({"type": "http", "headers": []})
+
+ @pytest.mark.asyncio
+ async def test_detach_skipped_on_generator_exit(self):
+ """A force-closed generator (client disconnect) must NOT call detach.
+
+ Regression for the "Token was created in a different Context" error:
when the
+ dependency generator is finalised via aclose() from a different
asyncio Task,
+ detaching the token would raise and OTel would log it at ERROR. We
skip detach
+ on the GeneratorExit unwind path instead.
+ """
+ with mock.patch.object(otel_context, "detach") as detach_spy:
+ gen = _extract_w3c_trace_context(self._make_request(),
dependency_solver="fastapi")
+ await gen.asend(None) # run up to and including the yield (after
attach)
+ await gen.aclose() # raises GeneratorExit at the yield
Review Comment:
This closes the generator from the same task that ran `asend`, so it doesn't
actually reproduce the original error. `cv.reset(token)` only raises "created
in a different Context" when the close runs in a *different* task; from the
same task it succeeds. The `detach` spy still guards the new behavior, so the
test isn't wrong, just narrower than the docstring implies. Separately, with
the route-error path no longer detaching (see my comment in app.py), the two
existing tests above (`test_trace_context_dep_cleans_up_on_route_exception` and
`test_route_exception_not_masked_by_detach_error`) now pass vacuously: the
first never asserts detach ran, and the second mocks `detach` to raise but
detach is never called on that path, so its side-effect never fires. Worth a
test that `athrow`s into the generator and asserts detach *is* called on the
route-error path.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
