auth: Use async version of getRequestMetadata

This avoids using DelayedStream and a thread hop when the credentials
are known immediately.

3 files changed, 86 insertions, 73 deletions

diff --git a/auth/build.gradle b/auth/build.gradle
index af50c28dd..151ec167f 100644
--- a/auth/build.gradle
+++ b/auth/build.gradle
@@ -2,6 +2,7 @@ description = "gRpc: Auth"
 dependencies {
     compile project(':grpc-core'),
             libraries.google_auth_credentials
-    testCompile libraries.oauth_client
+    testCompile project(':grpc-testing'),
+                libraries.oauth_client
     signature "org.codehaus.mojo.signature:java16:1.1@signature"
 }
diff --git a/auth/src/main/java/io/grpc/auth/GoogleAuthLibraryCallCredentials.java b/auth/src/main/java/io/grpc/auth/GoogleAuthLibraryCallCredentials.java
index d8035d3ac..cbfb1f6cc 100644
--- a/auth/src/main/java/io/grpc/auth/GoogleAuthLibraryCallCredentials.java
+++ b/auth/src/main/java/io/grpc/auth/GoogleAuthLibraryCallCredentials.java
@@ -19,6 +19,7 @@ package io.grpc.auth;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.auth.Credentials;
+import com.google.auth.RequestMetadataCallback;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.io.BaseEncoding;
 import io.grpc.Attributes;
@@ -84,48 +85,48 @@ final class GoogleAuthLibraryCallCredentials implements CallCredentials {
       applier.fail(e.getStatus());
       return;
     }
-    appExecutor.execute(new Runnable() {
-        @Override
-        public void run() {
-          try {
-            // Credentials is expected to manage caching internally if the metadata is fetched over
-            // the network.
-            //
-            // TODO(zhangkun83): we don't know whether there is valid cache data. If there is, we
-            // would waste a context switch by always scheduling in executor. However, we have to
-            // do so because we can't risk blocking the network thread. This can be resolved after
-            // https://github.com/google/google-auth-library-java/issues/3 is resolved.
-            //
-            // Some implementations may return null here.
-            Map<String, List<String>> metadata;
-            try {
-              metadata = creds.getRequestMetadata(uri);
-            } catch (IOException e) {
-              // Since it's an I/O failure, let the call be retried with UNAVAILABLE.
-              applier.fail(Status.UNAVAILABLE
-                  .withDescription("Credentials failed to obtain metadata")
-                  .withCause(e));
-              return;
-            }
-            // Re-use the headers if getRequestMetadata() returns the same map. It may return a
-            // different map based on the provided URI, i.e., for JWT. However, today it does not
-            // cache JWT and so we won't bother tring to save its return value based on the URI.
-            Metadata headers;
-            synchronized (GoogleAuthLibraryCallCredentials.this) {
-              if (lastMetadata == null || lastMetadata != metadata) {
-                lastMetadata = metadata;
-                lastHeaders = toHeaders(metadata);
-              }
-              headers = lastHeaders;
+    // Credentials is expected to manage caching internally if the metadata is fetched over
+    // the network.
+    creds.getRequestMetadata(uri, appExecutor, new RequestMetadataCallback() {
+      @Override
+      public void onSuccess(Map<String, List<String>> metadata) {
+        // Some implementations may pass null metadata.
+
+        // Re-use the headers if getRequestMetadata() returns the same map. It may return a
+        // different map based on the provided URI, i.e., for JWT. However, today it does not
+        // cache JWT and so we won't bother tring to save its return value based on the URI.
+        Metadata headers;
+        try {
+          synchronized (GoogleAuthLibraryCallCredentials.this) {
+            if (lastMetadata == null || lastMetadata != metadata) {
+              lastHeaders = toHeaders(metadata);
+              lastMetadata = metadata;
             }
-            applier.apply(headers);
-          } catch (Throwable e) {
-            applier.fail(Status.UNAUTHENTICATED
-                .withDescription("Failed computing credential metadata")
-                .withCause(e));
+            headers = lastHeaders;
           }
+        } catch (Throwable t) {
+          applier.fail(Status.UNAUTHENTICATED
+              .withDescription("Failed to convert credential metadata")
+              .withCause(t));
+          return;
         }
-      });
+        applier.apply(headers);
+      }
+
+      @Override
+      public void onFailure(Throwable e) {
+        if (e instanceof IOException) {
+          // Since it's an I/O failure, let the call be retried with UNAVAILABLE.
+          applier.fail(Status.UNAVAILABLE
+              .withDescription("Credentials failed to obtain metadata")
+              .withCause(e));
+        } else {
+          applier.fail(Status.UNAUTHENTICATED
+              .withDescription("Failed computing credential metadata")
+              .withCause(e));
+        }
+      }
+    });
   }
 
   /**
diff --git a/auth/src/test/java/io/grpc/auth/GoogleAuthLibraryCallCredentialsTest.java b/auth/src/test/java/io/grpc/auth/GoogleAuthLibraryCallCredentialsTest.java
index 1bf595b1e..04a48c07f 100644
--- a/auth/src/test/java/io/grpc/auth/GoogleAuthLibraryCallCredentialsTest.java
+++ b/auth/src/test/java/io/grpc/auth/GoogleAuthLibraryCallCredentialsTest.java
@@ -29,6 +29,7 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import com.google.auth.Credentials;
+import com.google.auth.RequestMetadataCallback;
 import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.OAuth2Credentials;
 import com.google.auth.oauth2.ServiceAccountCredentials;
@@ -43,6 +44,7 @@ import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
 import io.grpc.SecurityLevel;
 import io.grpc.Status;
+import io.grpc.testing.TestMethodDescriptors;
 import java.io.IOException;
 import java.net.URI;
 import java.security.KeyPair;
@@ -51,7 +53,9 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.Executor;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -75,19 +79,16 @@ public class GoogleAuthLibraryCallCredentialsTest {
       "Extra-Authorization-bin", Metadata.BINARY_BYTE_MARSHALLER);
 
   @Mock
-  private MethodDescriptor.Marshaller<String> stringMarshaller;
-
-  @Mock
-  private MethodDescriptor.Marshaller<Integer> intMarshaller;
-
-  @Mock
   private Credentials credentials;
 
   @Mock
   private MetadataApplier applier;
 
-  @Mock
-  private Executor executor;
+  private Executor executor = new Executor() {
+    @Override public void execute(Runnable r) {
+      pendingRunnables.add(r);
+    }
+  };
 
   @Captor
   private ArgumentCaptor<Metadata> headersCaptor;
@@ -95,8 +96,13 @@ public class GoogleAuthLibraryCallCredentialsTest {
   @Captor
   private ArgumentCaptor<Status> statusCaptor;
 
-  private MethodDescriptor<String, Integer> method;
-  private URI expectedUri;
+  private MethodDescriptor<Void, Void> method = MethodDescriptor.<Void, Void>newBuilder()
+      .setType(MethodDescriptor.MethodType.UNKNOWN)
+      .setFullMethodName("a.service/method")
+      .setRequestMarshaller(TestMethodDescriptors.voidMarshaller())
+      .setResponseMarshaller(TestMethodDescriptors.voidMarshaller())
+      .build();
+  private URI expectedUri = URI.create("https://testauthority/a.service");
 
   private final String authority = "testauthority";
   private final Attributes attrs = Attributes.newBuilder()
@@ -109,21 +115,32 @@ public class GoogleAuthLibraryCallCredentialsTest {
   @Before
   public void setUp() throws Exception {
     MockitoAnnotations.initMocks(this);
-    method = MethodDescriptor.<String, Integer>newBuilder()
-        .setType(MethodDescriptor.MethodType.UNKNOWN)
-        .setFullMethodName("a.service/method")
-        .setRequestMarshaller(stringMarshaller)
-        .setResponseMarshaller(intMarshaller)
-        .build();
-    expectedUri = new URI("https://testauthority/a.service");
     doAnswer(new Answer<Void>() {
-        @Override
-        public Void answer(InvocationOnMock invocation) {
-          Runnable r = (Runnable) invocation.getArguments()[0];
-          pendingRunnables.add(r);
+      @Override
+      public Void answer(InvocationOnMock invocation) {
+        Credentials mock = (Credentials) invocation.getMock();
+        URI uri = (URI) invocation.getArguments()[0];
+        RequestMetadataCallback callback = (RequestMetadataCallback) invocation.getArguments()[2];
+        Map<String, List<String>> metadata;
+        try {
+          // Default to calling the blocking method, since it is easier to mock
+          metadata = mock.getRequestMetadata(uri);
+        } catch (Exception ex) {
+          callback.onFailure(ex);
           return null;
         }
-      }).when(executor).execute(any(Runnable.class));
+        callback.onSuccess(metadata);
+        return null;
+      }
+    }).when(credentials).getRequestMetadata(
+        any(URI.class),
+        any(Executor.class),
+        any(RequestMetadataCallback.class));
+  }
+
+  @After
+  public void tearDown() {
+    assertEquals(0, pendingRunnables.size());
   }
 
   @Test
@@ -138,7 +155,6 @@ public class GoogleAuthLibraryCallCredentialsTest {
     GoogleAuthLibraryCallCredentials callCredentials =
         new GoogleAuthLibraryCallCredentials(credentials);
     callCredentials.applyRequestMetadata(method, attrs, executor, applier);
-    assertEquals(1, runPendingRunnables());
 
     verify(credentials).getRequestMetadata(eq(expectedUri));
     verify(applier).apply(headersCaptor.capture());
@@ -161,7 +177,6 @@ public class GoogleAuthLibraryCallCredentialsTest {
     GoogleAuthLibraryCallCredentials callCredentials =
         new GoogleAuthLibraryCallCredentials(credentials);
     callCredentials.applyRequestMetadata(method, attrs, executor, applier);
-    assertEquals(1, runPendingRunnables());
 
     verify(credentials).getRequestMetadata(eq(expectedUri));
     verify(applier).fail(statusCaptor.capture());
@@ -171,14 +186,13 @@ public class GoogleAuthLibraryCallCredentialsTest {
   }
 
   @Test
-  public void credentialsThrowsIoException() throws Exception {
+  public void credentialsFailsWithIoException() throws Exception {
     Exception exception = new IOException("Broken");
     when(credentials.getRequestMetadata(eq(expectedUri))).thenThrow(exception);
 
     GoogleAuthLibraryCallCredentials callCredentials =
         new GoogleAuthLibraryCallCredentials(credentials);
     callCredentials.applyRequestMetadata(method, attrs, executor, applier);
-    assertEquals(1, runPendingRunnables());
 
     verify(credentials).getRequestMetadata(eq(expectedUri));
     verify(applier).fail(statusCaptor.capture());
@@ -188,14 +202,13 @@ public class GoogleAuthLibraryCallCredentialsTest {
   }
 
   @Test
-  public void credentialsThrowsRuntimeException() throws Exception {
+  public void credentialsFailsWithRuntimeException() throws Exception {
     Exception exception = new RuntimeException("Broken");
     when(credentials.getRequestMetadata(eq(expectedUri))).thenThrow(exception);
 
     GoogleAuthLibraryCallCredentials callCredentials =
         new GoogleAuthLibraryCallCredentials(credentials);
     callCredentials.applyRequestMetadata(method, attrs, executor, applier);
-    assertEquals(1, runPendingRunnables());
 
     verify(credentials).getRequestMetadata(eq(expectedUri));
     verify(applier).fail(statusCaptor.capture());
@@ -218,7 +231,6 @@ public class GoogleAuthLibraryCallCredentialsTest {
       callCredentials.applyRequestMetadata(method, attrs, executor, applier);
     }
 
-    assertEquals(3, runPendingRunnables());
     verify(credentials, times(3)).getRequestMetadata(eq(expectedUri));
 
     verify(applier, times(3)).apply(headersCaptor.capture());
@@ -265,7 +277,6 @@ public class GoogleAuthLibraryCallCredentialsTest {
             .set(CallCredentials.ATTR_AUTHORITY, "example.com:443")
             .build(),
         executor, applier);
-    assertEquals(1, runPendingRunnables());
     verify(credentials).getRequestMetadata(eq(new URI("https://example.com/a.service")));
 
     callCredentials.applyRequestMetadata(method,
@@ -274,13 +285,13 @@ public class GoogleAuthLibraryCallCredentialsTest {
             .set(CallCredentials.ATTR_AUTHORITY, "example.com:123")
             .build(),
         executor, applier);
-    assertEquals(1, runPendingRunnables());
     verify(credentials).getRequestMetadata(eq(new URI("https://example.com:123/a.service")));
   }
 
   @Test
   public void serviceAccountToJwt() throws Exception {
     KeyPair pair = KeyPairGenerator.getInstance("RSA").generateKeyPair();
+    @SuppressWarnings("deprecation")
     ServiceAccountCredentials credentials = new ServiceAccountCredentials(
         null, "email@example.com", pair.getPrivate(), null, null) {
       @Override
@@ -292,7 +303,7 @@ public class GoogleAuthLibraryCallCredentialsTest {
     GoogleAuthLibraryCallCredentials callCredentials =
         new GoogleAuthLibraryCallCredentials(credentials);
     callCredentials.applyRequestMetadata(method, attrs, executor, applier);
-    assertEquals(1, runPendingRunnables());
+    assertEquals(0, runPendingRunnables());
 
     verify(applier).apply(headersCaptor.capture());
     Metadata headers = headersCaptor.getValue();
@@ -307,6 +318,7 @@ public class GoogleAuthLibraryCallCredentialsTest {
   public void serviceAccountWithScopeNotToJwt() throws Exception {
     final AccessToken token = new AccessToken("allyourbase", new Date(Long.MAX_VALUE));
     KeyPair pair = KeyPairGenerator.getInstance("RSA").generateKeyPair();
+    @SuppressWarnings("deprecation")
     ServiceAccountCredentials credentials = new ServiceAccountCredentials(
         null, "email@example.com", pair.getPrivate(), null, Arrays.asList("somescope")) {
       @Override
@@ -337,7 +349,6 @@ public class GoogleAuthLibraryCallCredentialsTest {
     GoogleAuthLibraryCallCredentials callCredentials =
         new GoogleAuthLibraryCallCredentials(credentials, null);
     callCredentials.applyRequestMetadata(method, attrs, executor, applier);
-    assertEquals(1, runPendingRunnables());
 
     verify(credentials).getRequestMetadata(eq(expectedUri));
     verify(applier).apply(headersCaptor.capture());