Implement JWT claims caching to avoid re-parsing on every #537
Conversation
- Added `JwtClaimsCache` as a shared state. - Updated `validate_request` to use JWT caching. - Updated error types to be `Clone` to allow caching to work.
Summary of ChangesHello @kamilkisiela, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a crucial performance optimization by implementing a caching mechanism for JSON Web Token (JWT) claims. By storing decoded and verified JWT payloads in a short-lived cache, the system avoids redundant processing for frequently accessed tokens, leading to a substantial increase in throughput. The changes involve integrating this cache into the shared application state, updating the JWT validation flow to leverage the cache, and making necessary adjustments to error types to support the new caching architecture. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a caching layer for JWT claims to improve performance, which is a great optimization. The changes involve adding a moka cache to the shared state and updating the JWT validation logic to use it.
My review has identified a critical issue in the implementation of validate_request that will prevent compilation and is also inefficient. I've also found a high-severity issue related to potential unbounded memory usage in the new cache, which violates the repository's style guide.
Please address these points to ensure the stability and performance of the router.
| .try_get_with(token.clone(), async { | ||
| let valid_jwks = self.jwks.all(); | ||
| self.authenticate(&valid_jwks, request) | ||
| .map(|(payload, _, _)| Arc::new(payload)) | ||
| }) |
There was a problem hiding this comment.
This implementation has a critical issue: it attempts to pass &mut HttpRequest into an async block. HttpRequest is not Send, so this will cause a compilation error because the request cannot be safely sent across threads.
Additionally, the current approach is inefficient. validate_request calls self.lookup(request) to get the token, and then the async block calls self.authenticate(...), which in turn calls self.lookup(request) again, performing the same work twice on a cache miss.
To fix this, the authentication logic should be performed directly on the token string that has already been extracted. This avoids the compilation error and the redundant work.
| .try_get_with(token.clone(), async { | |
| let valid_jwks = self.jwks.all(); | |
| self.authenticate(&valid_jwks, request) | |
| .map(|(payload, _, _)| Arc::new(payload)) | |
| }) | |
| .try_get_with(token.clone(), async { | |
| let valid_jwks = self.jwks.all(); | |
| let header = decode_header(&token).map_err(JwtError::InvalidJwtHeader)?; | |
| let jwk = self.find_matching_jwks(&header, &valid_jwks)?; | |
| let payload = self.decode_and_validate_token(&token, &jwk.keys)?; | |
| Ok(Arc::new(payload)) | |
| }) |
✅
|
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
|
🐋 This PR was built and pushed to the following Docker images: Image Names: Platforms: Image Tags: Docker metadata{
"buildx.build.ref": "builder-c9ea2d1f-37ee-4515-823c-c05d64739df0/builder-c9ea2d1f-37ee-4515-823c-c05d64739df00/p9oxvxmr3yf5fnnktdyo2fvpl",
"containerimage.descriptor": {
"mediaType": "application/vnd.oci.image.index.v1+json",
"digest": "sha256:b4d544aba0ea0cbcf932a15345c4831c2689481b26b725f2e2dcc3a6cebb4923",
"size": 1609
},
"containerimage.digest": "sha256:b4d544aba0ea0cbcf932a15345c4831c2689481b26b725f2e2dcc3a6cebb4923",
"image.name": "ghcr.io/graphql-hive/router:pr-537,ghcr.io/graphql-hive/router:sha-207f8c9"
} |
dotansimha
force-pushed
the
main
branch
4 times, most recently
from
4cffb36 to
c709bec
Compare
1 participant
Adds short-lived (5s) cache of decoded and verified JWT payloads.
With this optimization we go from 21.415 rps to 37.638 rps.
The 5s TTL is imo a good balance between performance and any cache invalidations that may occur.
Yeah I thought about doing it configurable, but does it really have to be? People may abuse it and set it to a high number that could interfere with expiration date of the token itself. We can always add it, when requested and needed.
JwtClaimsCacheas a shared state.validate_requestto use JWT caching.Cloneto allow caching to work.