RehanAhmad25 commented on code in PR #76467:
URL: https://github.com/apache/airflow/pull/76467#discussion_r4521927509
##########
airflow-core/src/airflow/api_fastapi/auth/tokens.py:
##########
@@ -129,18 +157,70 @@ def from_private_key(cls, *keys: AllowedPrivateKeys |
tuple[AllowedPrivateKeys,
return obj
async def fetch_jwks(self) -> None:
- if not self._should_fetch_jwks():
+ await self._load_if_needed(self._should_fetch_jwks)
+
+ async def _load_if_needed(self, needs_load: Callable[[], bool]) -> None:
+ """
+ Join the load in flight, or start one if ``needs_load`` holds.
+
+ A single load is shared by every event loop and thread, since
``validated_claims()`` runs
+ ``avalidated_claims()`` through ``async_to_sync`` on loops other than
the API server's.
+ """
+ load, started_here = self._join_or_start_load(needs_load)
+
+ if load is None:
return
- if self.url.startswith("http"):
- data = await self._fetch_remote_jwks()
- else:
- data = self._fetch_local_jwks()
+
+ if not started_here:
+ await asyncio.wrap_future(load)
+ return
+
+ try:
+ await self._load_jwks()
+ finally:
+ self._finish_load(load)
Review Comment:
`set_running_or_notify_cancel()` protects the shared future if a waiting
request is cancelled, but what happens if the request that starts
`_load_jwks()` is cancelled while the fetch is in progress? The `finally` block
still resolves the shared future as successful, even though the keyset may not
have been updated.
Waiting requests then check the old keyset and fail validation with `Kid did
not match any validation keys`. Since `last_fetch_attempt_at` was already
updated, the retry interval also prevents another immediate reload.
Could we keep the load running if its initiating request is cancelled, or
allow a retry or handoff, and cover this case in a test?
--
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]
