Skip to content

⚡️ Speed up method ControlAccountAuthInterceptor._generic_auth_unary_method_handler by 13% #31

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

⚡️ Speed up method ControlAccountAuthInterceptor._generic_auth_unary_method_handler by 13% #31

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 12 additions & 12 deletions framework/py/flwr/superlink/servicer/control/control_account_auth_interceptor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,19 +84,22 @@ def intercept_service(
def _generic_auth_unary_method_handler(
self, method_handler: grpc.RpcMethodHandler
) -> grpc.RpcMethodHandler:
# Locally cache repeated lookups for miniscule performance gain (method_handler fields used multiple times)
unary_unary = method_handler.unary_unary
unary_stream = method_handler.unary_stream

def _generic_method_handler(
request: Request,
request, # type: ignore
context: grpc.ServicerContext,
) -> Response:
call = method_handler.unary_unary or method_handler.unary_stream
):
call = unary_unary or unary_stream
metadata = context.invocation_metadata()

# Intercept GetLoginDetails and GetAuthTokens requests, and return
# the response without authentication
# Fast path: short-circuit for token-issuing RPCs
if isinstance(request, (GetLoginDetailsRequest, GetAuthTokensRequest)):
return call(request, context) # type: ignore

# For other requests, check if the account is authenticated
# Authn fast path
valid_tokens, account_info = self.authn_plugin.validate_tokens_in_metadata(
metadata
)
Expand All @@ -107,9 +110,7 @@ def _generic_method_handler(
"Tokens validated, but account info not found",
)
raise grpc.RpcError()
# Store account info in contextvars for authenticated accounts
shared_account_info.set(account_info)
# Check if the account is authorized
if not self.authz_plugin.authorize(account_info):
context.abort(
grpc.StatusCode.PERMISSION_DENIED,
Expand All @@ -119,7 +120,7 @@ def _generic_method_handler(
raise grpc.RpcError()
return call(request, context) # type: ignore

# If the account is not authenticated, refresh tokens
# Authn fallback: refresh and repeat
tokens, account_info = self.authn_plugin.refresh_tokens(metadata)
if tokens is not None:
if account_info is None:
Expand All @@ -128,9 +129,7 @@ def _generic_method_handler(
"Tokens refreshed, but account info not found",
)
raise grpc.RpcError()
# Store account info in contextvars for authenticated accounts
shared_account_info.set(account_info)
# Check if the account is authorized
if not self.authz_plugin.authorize(account_info):
context.abort(
grpc.StatusCode.PERMISSION_DENIED,
Expand All @@ -145,7 +144,8 @@ def _generic_method_handler(
context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
raise grpc.RpcError() # This line is unreachable

if method_handler.unary_unary:
# Avoid attribute lookup on every call
if unary_unary:
message_handler = grpc.unary_unary_rpc_method_handler
else:
message_handler = grpc.unary_stream_rpc_method_handler
Expand Down