From af4a3a064e6974099dd5333f291ac7417bc57978 Mon Sep 17 00:00:00 2001 From: Mateus Azis Date: Wed, 20 Dec 2023 11:28:18 -0800 Subject: [PATCH] Using future callback to remove cancelled/failed auth future from map. --- .../internal/BinderTransportSecurity.java | 77 ++++++++++--------- 1 file changed, 39 insertions(+), 38 deletions(-) diff --git a/binder/src/main/java/io/grpc/binder/internal/BinderTransportSecurity.java b/binder/src/main/java/io/grpc/binder/internal/BinderTransportSecurity.java index e210e0440b7b..a6245da89c22 100644 --- a/binder/src/main/java/io/grpc/binder/internal/BinderTransportSecurity.java +++ b/binder/src/main/java/io/grpc/binder/internal/BinderTransportSecurity.java @@ -19,6 +19,7 @@ import com.google.common.util.concurrent.FutureCallback; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.MoreExecutors; import io.grpc.Attributes; import io.grpc.Internal; @@ -32,12 +33,10 @@ import io.grpc.Status; import io.grpc.internal.GrpcAttributes; +import java.util.HashMap; import java.util.concurrent.CancellationException; -import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; -import java.util.concurrent.Future; - import javax.annotation.CheckReturnValue; import javax.annotation.Nullable; @@ -106,14 +105,14 @@ public ServerCall.Listener interceptCall( // Most SecurityPolicy will have synchronous implementations that provide an // immediately-resolved Future. In that case, short-circuit to avoid unnecessary allocations // and asynchronous code if the authorization result is already present. - if (!authStatusFuture.isDone() || authStatusFuture.isCancelled()) { + if (!authStatusFuture.isDone()) { return newServerCallListenerForPendingAuthResult(authStatusFuture, call, headers, next); } Status authStatus; try { authStatus = Futures.getDone(authStatusFuture); - } catch (ExecutionException e) { + } catch (ExecutionException | CancellationException e) { // Failed futures are treated as an internal error rather than a security rejection. authStatus = Status.INTERNAL.withCause(e); } @@ -163,12 +162,12 @@ public void onFailure(Throwable t) { private static final class TransportAuthorizationState { private final int uid; private final ServerPolicyChecker serverPolicyChecker; - private final ConcurrentHashMap> serviceAuthorization; + private final HashMap> serviceAuthorization; TransportAuthorizationState(int uid, ServerPolicyChecker serverPolicyChecker) { this.uid = uid; this.serverPolicyChecker = serverPolicyChecker; - serviceAuthorization = new ConcurrentHashMap<>(8); + serviceAuthorization = new HashMap<>(8); } /** Get whether we're authorized to make this call. */ @@ -179,40 +178,42 @@ ListenableFuture checkAuthorization(MethodDescriptor method) { // which is true for all generated methods. Otherwise, programmatically // created methods could cause this cache to grow unbounded. boolean useCache = method.isSampledToLocalTracing(); - if (useCache) { - @Nullable ListenableFuture authorization = serviceAuthorization.get(serviceName); - if (authorization != null && !isFailedOrCancelled(authorization)) { - // Authorization check exists and is a pending or successful future (even if for a failed - // authorization). - return authorization; + + @Nullable ListenableFuture authorization; + synchronized (serviceAuthorization) { + if (useCache) { + authorization = serviceAuthorization.get(serviceName); + if (authorization != null) { + // Authorization check exists and is a pending or successful future (even if for a + // failed authorization). + return authorization; + } + } + + + // Under high load, this may trigger a large number of concurrent authorization checks that + // perform essentially the same work and have the potential of exhausting the resources they + // depend on. This was a non-issue in the past with synchronous policy checks due to the + // fixed-size nature of the thread pool this method runs under. + // + // TODO(10669): evaluate if there should be at most a single pending authorization check per + // (uid, serviceName) pair at any given time. + authorization = serverPolicyChecker.checkAuthorizationForServiceAsync(uid, serviceName); + if (useCache) { + serviceAuthorization.put(serviceName, authorization); } - serviceAuthorization.remove(serviceName); - } - // Under high load, this may trigger a large number of concurrent authorization checks that - // perform essentially the same work and have the potential of exhausting the resources they - // depend on. This was a non-issue in the past with synchronous policy checks due to the - // fixed-size nature of the thread pool this method runs under. - // - // TODO(10669): evaluate if there should be at most a single pending authorization check per - // (uid, serviceName) pair at any given time. - ListenableFuture authorization = - serverPolicyChecker.checkAuthorizationForServiceAsync(uid, serviceName); - if (useCache) { - serviceAuthorization.putIfAbsent(serviceName, authorization); } - return authorization; - } - } - private static boolean isFailedOrCancelled(Future doneFuture) { - if (!doneFuture.isDone()) { - return false; - } - try { - T unused = Futures.getDone(doneFuture); - return false; - } catch (ExecutionException | CancellationException e) { - return true; + final ListenableFuture finalAuthorization = authorization; + + return Futures.catching(authorization, Throwable.class, t -> { + synchronized (serviceAuthorization) { + if (useCache) { + serviceAuthorization.remove(serviceName, finalAuthorization); + } + } + throw new IllegalStateException(t); + }, MoreExecutors.directExecutor()); } }