Repository: GoogleCloudPlatform/cloud-bigtable-client (java)

Issue:
RefreshingOAuth2CredentialsInterceptor should respect deadlines
Currently the RefreshingOAuth2CredentialsInterceptor is hardcoded to wait up to 15 seconds for a background thread to refresh the token. Instead it should check if a user provided deadline is given in the CallOptions and use that if it's earlier than the hardcoded one

Pull request description:
RefreshingOAuth2CredentialsInterceptor now honors CallOptions Deadlines
This should fix #1718.

Patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/io/OAuthCredentialsCache.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/io/OAuthCredentialsCache.java
index 8231f4287..f2b2cec75 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/io/OAuthCredentialsCache.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/io/OAuthCredentialsCache.java
@@ -75,8 +75,6 @@ public class OAuthCredentialsCache {
     }
   }
 
-  // TODO: this should split into cache type operations and a token used by RefreshingOAuth2CredentialsInterceptor
-
   static class HeaderToken {
     private final Status status;
     private final String header;
@@ -183,9 +181,9 @@ public class OAuthCredentialsCache {
     return headerCache;
   }
 
-  HeaderToken getHeader(int timeoutSeconds) {
+  HeaderToken getHeader(long timeoutMs) {
     try {
-      return getHeaderUnsafe(timeoutSeconds).getToken();
+      return getHeaderUnsafe(timeoutMs).getToken();
     } catch (Exception e) {
       LOG.warn("Got an unexpected exception while trying to refresh google credentials.", e);
       Status status = Status.UNAUTHENTICATED
@@ -197,9 +195,9 @@ public class OAuthCredentialsCache {
 
   /**
    * Get the http credential header we need from a new oauth2 AccessToken.
-   * @param timeoutSeconds
+   * @param timeoutMs
    */
-  private HeaderCacheElement getHeaderUnsafe(int timeoutSeconds) {
+  private HeaderCacheElement getHeaderUnsafe(long timeoutMs) {
     // Optimize for the common case: do a volatile read to peek for a Good cache value
     HeaderCacheElement headerCacheUnsync = this.headerCache;
 
@@ -213,7 +211,7 @@ public class OAuthCredentialsCache {
       case Expired:
       case Exception:
         // defer the future resolution (asyncRefresh will spin up a thread that will try to acquire the lock)
-        return syncRefresh(timeoutSeconds);
+        return syncRefresh(timeoutMs);
       default:
         String message = "Could not process state: " + headerCacheUnsync.getCacheState();
         LOG.warn(message);
@@ -228,9 +226,9 @@ public class OAuthCredentialsCache {
    * been refreshed.
    * This method should not be called while holding the refresh lock
    */
-  private HeaderCacheElement syncRefresh(int timeoutSeconds) {
+  private HeaderCacheElement syncRefresh(long timeoutMs) {
     try (Closeable ss = Tracing.getTracer().spanBuilder("CredentialsRefresh").startScopedSpan()) {
-      return asyncRefresh().get(timeoutSeconds, TimeUnit.SECONDS);
+      return asyncRefresh().get(timeoutMs, TimeUnit.MILLISECONDS);
     } catch (InterruptedException e) {
       LOG.warn("Interrupted while trying to refresh google credentials.", e);
       Thread.currentThread().interrupt();
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 699b6b5a0..7b7817983 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
@@ -32,6 +32,7 @@ import io.grpc.MethodDescriptor;
 import io.grpc.Status;
 
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Client interceptor that authenticates all calls by binding header data provided by a credential.
@@ -107,7 +108,7 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
    */
   @Override
   public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method,
-      CallOptions callOptions, Channel next) {
+      final CallOptions callOptions, Channel next) {
     return new SimpleForwardingClientCall<ReqT, RespT>(next.newCall(method, callOptions)) {
 
       /**
@@ -119,7 +120,13 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
 
       @Override
       public void start(Listener<RespT> responseListener, Metadata headers) {
-        HeaderToken token = store.getHeader(TIMEOUT_SECONDS);
+        long timeoutMs = -1;
+        if (callOptions.getDeadline() != null) {
+          timeoutMs = callOptions.getDeadline().timeRemaining(TimeUnit.MILLISECONDS);
+        } else {
+          timeoutMs = TimeUnit.SECONDS.toMillis(TIMEOUT_SECONDS);
+        }
+        HeaderToken token = store.getHeader(timeoutMs);
 
         if (!token.getStatus().isOk()) {
           unauthorized = true;


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 92b0a52df..29cf25ab1 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
@@ -15,8 +15,8 @@
  */
 package com.google.cloud.bigtable.grpc.io;
 
-import static org.mockito.Matchers.*;
 import static org.mockito.Mockito.*;
+import static org.junit.Assert.*;
 
 import com.google.bigtable.v2.BigtableGrpc;
 import com.google.bigtable.v2.ReadRowsRequest;
@@ -26,6 +26,7 @@ import com.google.auth.oauth2.OAuth2Credentials;
 import java.io.IOException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 
 import io.grpc.*;
 import org.junit.AfterClass;
@@ -36,6 +37,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.ArgumentCaptor;
+import org.mockito.Matchers;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -74,14 +76,14 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
     private ClientCall.Listener mockListener;
 
     @Mock
-    private OAuthCredentialsCache mockStore;
+    private OAuthCredentialsCache mockCache;
 
     @Before
     public void setupMocks() {
         MockitoAnnotations.initMocks(this);
         when(mockChannel.newCall(any(MethodDescriptor.class), any(CallOptions.class)))
                 .thenReturn(mockClientCall);
-        underTest = new RefreshingOAuth2CredentialsInterceptor(mockStore);
+        underTest = new RefreshingOAuth2CredentialsInterceptor(mockCache);
     }
 
     @Test
@@ -95,6 +97,42 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
         checkOKCompletedCorrectly();
     }
 
+    @Test
+    /**
+     * Ensures that default deadline is honored
+     */
+    public void testDefaultDeadline() throws IOException {
+        initializeOk();
+        sendRequest(CallOptions.DEFAULT);
+        assertEquals(RefreshingOAuth2CredentialsInterceptor.TIMEOUT_SECONDS * 1000, getDeadlineMs());
+    }
+
+    @Test
+    /**
+     * Ensures that short deadlines in CallOptions are honored
+     */
+    public void testShortDeadline() throws IOException {
+        initializeOk();
+        sendRequest(CallOptions.DEFAULT.withDeadlineAfter(2, TimeUnit.SECONDS));
+        assertTrue(getDeadlineMs() <= 2000);
+    }
+
+    @Test
+    /**
+     * Ensures that long deadlines in CallOptions are honored
+     */
+    public void testLongDeadline() throws IOException {
+        initializeOk();
+        sendRequest(CallOptions.DEFAULT.withDeadlineAfter(60, TimeUnit.SECONDS));
+        assertTrue(getDeadlineMs() >= 59000);
+    }
+
+    private long getDeadlineMs() {
+        ArgumentCaptor<Long> timeoutCaptor = ArgumentCaptor.forClass(Long.class);
+        verify(mockCache, times(1)).getHeader(timeoutCaptor.capture());
+        return timeoutCaptor.getValue();
+    }
+
     @Test
     /**
      * Test to make sure that the interceptor revokes credentials if an RPC receives an UNAUTHENTICATED status.
@@ -103,7 +141,7 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
         Status grpcStatus = Status.UNAUTHENTICATED;
 
         // Something bad happened, and authentication could not be established
-        when(mockStore.getHeader(anyInt()))
+        when(mockCache.getHeader(anyInt()))
                 .thenReturn(new OAuthCredentialsCache.HeaderToken(grpcStatus, null));
 
         ClientCall<ReadRowsRequest, ReadRowsResponse> call = underTest.interceptCall(
@@ -141,13 +179,13 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
         initializeOk();
 
         // 10 requests come in around the same time, and all of them get UNAUTHENTICATED
-        for (int i = 0; i < 10; i++){
+        for (int i = 0; i < 10; i++) {
             sendRequest(CallOptions.DEFAULT);
             getListener().onClose(Status.UNAUTHENTICATED, new Metadata());
         }
 
         // only a single revoke should be submitted due to rate limiting
-        verify(mockStore, times(1)).revokeUnauthToken(
+        verify(mockCache, times(1)).revokeUnauthToken(
                 any(OAuthCredentialsCache.HeaderToken.class));
     }
 
@@ -180,7 +218,7 @@ public class RefreshingOAuth2CredentialsInterceptorTest {
     }
 
     private void initializeOk() {
-        when(mockStore.getHeader(anyInt()))
+        when(mockCache.getHeader(anyInt()))
                 .thenReturn(new OAuthCredentialsCache.HeaderToken(Status.OK, HEADER));
     }