Repository: GoogleCloudPlatform/cloud-bigtable-client
Pull request #1657 (2018-02-08 18:35:18)

Auth timeout is now a variable and 15 seconds.


Patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/io/RefreshingOAuth2CredentialsInterceptor.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/io/RefreshingOAuth2CredentialsInterceptor.java
index 627e43863..2fc3cee90 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/io/RefreshingOAuth2CredentialsInterceptor.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/io/RefreshingOAuth2CredentialsInterceptor.java
@@ -59,6 +59,8 @@ import javax.annotation.concurrent.GuardedBy;
  */
 public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor {
 
+  public static int TIMEOUT_SECONDS = 15;
+
   private static final Logger LOG = new Logger(RefreshingOAuth2CredentialsInterceptor.class);
 
   /**
@@ -75,7 +77,7 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
   enum CacheState {
     Good, Stale, Expired, Exception
   }
-  
+
   private static final Metadata.Key<String> AUTHORIZATION_HEADER_KEY = Metadata.Key.of(
       "Authorization", Metadata.ASCII_STRING_MARSHALLER);
 
@@ -143,28 +145,18 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
   private static final HeaderCacheElement EMPTY_HEADER = new HeaderCacheElement(null, 0);
 
   private final ExecutorService executor;
-
-  @VisibleForTesting
-  RateLimiter rateLimiter;
-
+  private final RateLimiter rateLimiter;
   private final OAuth2Credentials credentials;
+  private final Object lock = new Object();
 
-  @VisibleForTesting
-  final Object lock = new Object();
+  @GuardedBy("lock")
+  private Future<HeaderCacheElement> futureToken = null;
 
   // Note that the cache is volatile to allow us to peek for a Good value
   @VisibleForTesting
   @GuardedBy("lock")
   volatile HeaderCacheElement headerCache = EMPTY_HEADER;
 
-  @VisibleForTesting
-  @GuardedBy("lock")
-  boolean isRefreshing = false;
-
-  @VisibleForTesting
-  @GuardedBy("lock")
-  Future<HeaderCacheElement> futureToken = null;
-
 
   /**
    * <p>Constructor for RefreshingOAuth2CredentialsInterceptor.</p>
@@ -288,7 +280,7 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
           );
       }
     }
-    return deferredResult.get(5, TimeUnit.SECONDS);
+    return deferredResult.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
   }
 
   /**
@@ -298,7 +290,7 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
    */
   HeaderCacheElement syncRefresh() {
     try {
-      return asyncRefresh().get(5, TimeUnit.SECONDS);
+      return asyncRefresh().get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
     } catch (Exception e) {
       return new HeaderCacheElement(
           Status.UNAUTHENTICATED
@@ -315,27 +307,31 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
     LOG.trace("asyncRefresh");
 
     synchronized (lock) {
-      if (isRefreshing) {
+      if (futureToken != null) {
         LOG.trace("asyncRefresh is already in progress");
         return futureToken;
       }
-      isRefreshing = true;
       LOG.trace("asyncRefresh taking ownership");
 
+      Future<HeaderCacheElement> future;
       try {
-        this.futureToken = executor.submit(new Callable<HeaderCacheElement>() {
+        future = executor.submit(new Callable<HeaderCacheElement>() {
           @Override
           public HeaderCacheElement call() throws Exception {
             HeaderCacheElement newToken = refreshCredentials();
             return updateToken(newToken);
           }
         });
+
+        if (!future.isDone()) {
+          this.futureToken = future;
+        }
       } catch (RuntimeException e) {
-        isRefreshing = false;
+        futureToken = null;
         throw e;
       }
 
-      return this.futureToken;
+      return future;
     }
   }
 
@@ -354,7 +350,6 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
             + "New token state: {}, status: {}", newState, newToken.status);
       }
       futureToken = null;
-      isRefreshing = false;
 
       return headerCache;
     }
@@ -413,4 +408,11 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
       super.onClose(status, trailers);
     }
   }
+
+  @VisibleForTesting
+  boolean isRefreshing() {
+    synchronized(lock) {
+      return futureToken != null;
+    }
+  }
 }


Test patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/io/RefreshingOAuth2CredentialsInterceptorTest.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/io/RefreshingOAuth2CredentialsInterceptorTest.java
index fe5d4273c..19196bca1 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/io/RefreshingOAuth2CredentialsInterceptorTest.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/io/RefreshingOAuth2CredentialsInterceptorTest.java
@@ -19,7 +19,7 @@ import static org.hamcrest.CoreMatchers.containsString;
 import static org.mockito.Mockito.times;
 
 import com.google.api.client.util.Clock;
-import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.OAuth2Credentials;
 import com.google.cloud.bigtable.grpc.io.RefreshingOAuth2CredentialsInterceptor.CacheState;
@@ -86,6 +86,18 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
   public void testSyncRefresh() throws IOException {
     initialize(HeaderCacheElement.TOKEN_STALENESS_MS + 1);
     Assert.assertEquals(CacheState.Good, underTest.headerCache.getCacheState());
+    Assert.assertFalse(underTest.isRefreshing());
+  }
+
+  @Test
+  public void testRefresh() throws IOException {
+    Mockito.when(credentials.refreshAccessToken()).thenReturn(
+        new AccessToken("", new Date(HeaderCacheElement.TOKEN_STALENESS_MS + 1)));
+    underTest = new RefreshingOAuth2CredentialsInterceptor(MoreExecutors.newDirectExecutorService(),
+        credentials);
+    underTest.syncRefresh();
+    Assert.assertEquals(CacheState.Good, underTest.headerCache.getCacheState());
+    Assert.assertFalse(underTest.isRefreshing());
   }
 
   @Test
@@ -116,7 +128,6 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
   @Test
   public void testRefreshAfterFailure() throws Exception {
     underTest = new RefreshingOAuth2CredentialsInterceptor(executorService, credentials);
-    underTest.rateLimiter.setRate(100000);
 
     final AccessToken accessToken = new AccessToken("hi", new Date(HeaderCacheElement.TOKEN_STALENESS_MS + 1));
 
@@ -142,7 +153,6 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
   @Test
   public void testRefreshAfterStale() throws Exception {
     underTest = new RefreshingOAuth2CredentialsInterceptor(executorService, credentials);
-    underTest.rateLimiter.setRate(100000);
 
     final AccessToken staleToken = new AccessToken("stale", new Date(HeaderCacheElement.TOKEN_STALENESS_MS + 1));
     AccessToken goodToken = new AccessToken("good", new Date(HeaderCacheElement.TOKEN_STALENESS_MS + 11));
@@ -168,11 +178,7 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
     Assert.assertThat(secondResult.header, containsString("stale"));
 
     // Wait for the refresh to finish
-    final Future<?> waiter;
-    synchronized (underTest.lock) {
-      waiter = underTest.isRefreshing ? underTest.futureToken : Futures.immediateFuture(null);
-    }
-    waiter.get();
+    underTest.syncRefresh();
 
     // Third call - now returns good token
     HeaderCacheElement thirdResult = underTest.getHeaderSafe();
@@ -228,7 +234,6 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
 
     underTest =
         new RefreshingOAuth2CredentialsInterceptor(executorService, credentials);
-    underTest.rateLimiter.setRate(100000);
 
     // At this point, the access token wasn't retrieved yet. The
     // RefreshingOAuth2CredentialsInterceptor considers null to be Expired.
@@ -240,9 +245,7 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
     Assert.assertEquals(CacheState.Stale, underTest.headerCache.getCacheState());
 
     // Check to make sure we're no longer refreshing.
-    synchronized (underTest.lock) {
-      Assert.assertFalse(underTest.isRefreshing);
-    }
+    Assert.assertFalse(underTest.isRefreshing());
 
     // Kick off a couple of asynchronous refreshes. Kicking off more than one shouldn't be
     // necessary, but also should not be harmful, since there are likely to be multiple concurrent
@@ -252,7 +255,7 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
     underTest.asyncRefresh();
 
     syncCall(lock, syncRefreshCallable);
-    Assert.assertFalse(underTest.isRefreshing);
+    Assert.assertFalse(underTest.isRefreshing());
   }
 
   private void syncCall(final Object lock, Callable<Void> syncRefreshCallable)
@@ -269,9 +272,7 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
     // actually doing a refresh at this point; the other ones will have see that a refresh is in
     // progress and finish the invocation of the Thread without performing a refres().. Make sure
     // that at least 1 refresh process is in progress.
-    synchronized (underTest.lock) {
-      Assert.assertTrue(underTest.isRefreshing);
-    }
+    Assert.assertTrue(underTest.isRefreshing());
 
     synchronized (lock) {
       lock.notifyAll();