Repository: GoogleCloudPlatform/cloud-bigtable-client
Pull request #1726 (2018-03-13 21:01:39)

   Breaking down RefreshingOAuth2CredentialsInterceptor into two classes
`RefreshingOAuth2CredentialsInterceptor` managed oauth2 tokens asynchronously and integrated with grpc.  There was a lot of code associated with `RefreshingOAuth2CredentialsInterceptor`, and it is difficult to tell if everything is properly tested.  Testing of the grpc specific integration was incomplete.  To fix this problem, this PR breaks down `RefreshingOAuth2CredentialsInterceptor` into two components:
    
1. `RefreshingOAuth2CredentialsInterceptor` will only handle the grpc aspects.  It will act as a ClientCall.Listener that is invoked by a gRPC call chain.  It will delegate the details of managing the auth header to OAuthCredentialsStore.  It will also manage the functionality of handling UNAUTHENTICATED RPCs, including throttling

 2. `OAuthCredentialsStore` will handle the lifecycle of an OAuth2 token
    
    There should be no public facing behavioral changes or interface changes.  `RefreshingOAuth2CredentialsInterceptor` maintains its `asyncRefresh()` method and the `TIMEOUT_SECONDS` public variable.
    
    Missing tests were introduced, and some of the tests were refactored, but remained the same functionally.

spollapally: LGTM

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
new file mode 100644
index 000000000..8231f4287
--- /dev/null
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/io/OAuthCredentialsCache.java
@@ -0,0 +1,357 @@
+/*
+ * Copyright 2018 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable.grpc.io;
+
+import com.google.api.client.util.Clock;
+import com.google.auth.oauth2.AccessToken;
+import com.google.auth.oauth2.OAuth2Credentials;
+import com.google.cloud.bigtable.config.Logger;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import com.google.common.util.concurrent.Futures;
+import io.opencensus.trace.Tracing;
+import io.grpc.Status;
+
+import java.io.Closeable;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import javax.annotation.concurrent.GuardedBy;
+
+/**
+ * This class caches calls to {@link OAuth2Credentials#refreshAccessToken()}.  It asynchronously refreshes
+ * the token when it becomes stale.
+ *
+ * @author sduskis
+ * @version $Id: $Id
+ */
+public class OAuthCredentialsCache {
+
+  private static final Logger LOG = new Logger(OAuthCredentialsCache.class);
+  private static final HeaderCacheElement EMPTY_HEADER = new HeaderCacheElement(null, 0);
+
+  @VisibleForTesting
+  static Clock clock = Clock.SYSTEM;
+
+  /**
+   * <p>
+   * This enum describes the states of the OAuth header.
+   * </p>
+   * <ol>
+   * <li>Good - fine to use, and does not need to be refreshed.
+   * <li>Stale - fine to use, but requires an async refresh
+   * <li>Expired - Cannot be used. Wait for a new token to be loaded
+   * </ol>
+   */
+  @VisibleForTesting
+  enum CacheState {
+    Good(true), Stale(true),
+    Expired(false), Exception(false);
+
+    private boolean isValid;
+
+    CacheState(boolean isValid) {
+      this.isValid = isValid;
+    }
+
+    public boolean isValid() {
+      return isValid;
+    }
+  }
+
+  // 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;
+
+    public HeaderToken(Status status, String header) {
+      this.status = status;
+      this.header = header;
+    }
+
+    Status getStatus() {
+      return status;
+    }
+
+    String getHeader() {
+      return header;
+    }
+  }
+
+  static class HeaderCacheElement {
+
+    /**
+     * This specifies how far in advance of a header expiration do we consider the token stale. The
+     * Stale state indicates that the interceptor needs to do an asynchronous refresh.
+     */
+    static final long TOKEN_STALENESS_MS = TimeUnit.MINUTES.toMillis(6);
+
+    /**
+     * After the token is "expired," the interceptor blocks gRPC calls. The Expired state indicates
+     * that the interceptor needs to do a synchronous refresh.
+     */
+    // 5 minutes as per https://github.com/google/google-auth-library-java/pull/95
+    static final long TOKEN_EXPIRES_MS = TimeUnit.MINUTES.toMillis(5);
+
+    private final HeaderToken token;
+    private final long actualExpirationTimeMs;
+
+    private HeaderCacheElement(AccessToken token) {
+      if (token.getExpirationTime() == null) {
+        actualExpirationTimeMs = Long.MAX_VALUE;
+      } else {
+        actualExpirationTimeMs = token.getExpirationTime().getTime();
+      }
+      this.token = new HeaderToken(Status.OK, "Bearer " + token.getTokenValue());
+    }
+
+    private HeaderCacheElement(String header, long actualExpirationTimeMs) {
+      this.token = new HeaderToken(Status.OK, header);
+      this.actualExpirationTimeMs = actualExpirationTimeMs;
+    }
+
+    private HeaderCacheElement(Status errorStatus) {
+      Preconditions.checkArgument(!errorStatus.isOk(), "Error status can't be OK");
+      this.token = new HeaderToken(errorStatus, null);
+      this.actualExpirationTimeMs = 0;
+    }
+
+    CacheState getCacheState() {
+      long now = clock.currentTimeMillis();
+
+      if (!token.status.isOk()) {
+        return CacheState.Exception;
+      } else if (actualExpirationTimeMs - TOKEN_EXPIRES_MS <= now) {
+        return CacheState.Expired;
+      } else if (actualExpirationTimeMs - TOKEN_STALENESS_MS <= now) {
+        return CacheState.Stale;
+      } else {
+        return CacheState.Good;
+      }
+    }
+
+    private boolean isValid() {
+      return getCacheState().isValid();
+    }
+
+    HeaderToken getToken() {
+      return token;
+    }
+  }
+
+  private final ExecutorService executor;
+  private final OAuth2Credentials credentials;
+  private 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
+  @GuardedBy("lock")
+  private volatile HeaderCacheElement headerCache = EMPTY_HEADER;
+
+  /**
+   * <p>Constructor for RefreshingOAuth2CredentialsInterceptor.</p>
+   *
+   * @param scheduler a {@link ExecutorService} object.
+   * @param credentials a {@link OAuth2Credentials} object.
+   */
+  public OAuthCredentialsCache(ExecutorService scheduler, OAuth2Credentials credentials) {
+    this.executor = Preconditions.checkNotNull(scheduler);
+    this.credentials = Preconditions.checkNotNull(credentials);
+  }
+
+  @VisibleForTesting
+  HeaderCacheElement getHeaderCache() {
+    return headerCache;
+  }
+
+  HeaderToken getHeader(int timeoutSeconds) {
+    try {
+      return getHeaderUnsafe(timeoutSeconds).getToken();
+    } catch (Exception e) {
+      LOG.warn("Got an unexpected exception while trying to refresh google credentials.", e);
+      Status status = Status.UNAUTHENTICATED
+              .withDescription("Unexpected failure get auth token")
+              .withCause(e);
+      return new HeaderToken(status, null);
+    }
+  }
+
+  /**
+   * Get the http credential header we need from a new oauth2 AccessToken.
+   * @param timeoutSeconds
+   */
+  private HeaderCacheElement getHeaderUnsafe(int timeoutSeconds) {
+    // Optimize for the common case: do a volatile read to peek for a Good cache value
+    HeaderCacheElement headerCacheUnsync = this.headerCache;
+
+    // TODO(igorbernstein2): figure out how to make this work with appengine request scoped threads
+    switch (headerCacheUnsync.getCacheState()) {
+      case Good:
+        return headerCacheUnsync;
+      case Stale:
+        asyncRefresh();
+        return headerCacheUnsync;
+      case Expired:
+      case Exception:
+        // defer the future resolution (asyncRefresh will spin up a thread that will try to acquire the lock)
+        return syncRefresh(timeoutSeconds);
+      default:
+        String message = "Could not process state: " + headerCacheUnsync.getCacheState();
+        LOG.warn(message);
+        return new HeaderCacheElement(
+            Status.UNAUTHENTICATED
+                .withCause(new IllegalStateException(message)));
+    }
+  }
+
+  /**
+   * Refresh the credentials and block. Will return an error if the credentials haven't
+   * been refreshed.
+   * This method should not be called while holding the refresh lock
+   */
+  private HeaderCacheElement syncRefresh(int timeoutSeconds) {
+    try (Closeable ss = Tracing.getTracer().spanBuilder("CredentialsRefresh").startScopedSpan()) {
+      return asyncRefresh().get(timeoutSeconds, TimeUnit.SECONDS);
+    } catch (InterruptedException e) {
+      LOG.warn("Interrupted while trying to refresh google credentials.", e);
+      Thread.currentThread().interrupt();
+      return new HeaderCacheElement(
+          Status.UNAUTHENTICATED
+              .withDescription("Authentication was interrupted.")
+              .withCause(e)
+      );
+    } catch (ExecutionException e) {
+      LOG.warn("ExecutionException while trying to refresh google credentials.", e);
+      return new HeaderCacheElement(
+          Status.UNAUTHENTICATED
+              .withDescription("ExecutionException during Authentication.")
+              .withCause(e)
+      );
+    } catch (TimeoutException e) {
+      LOG.warn("TimeoutException while trying to refresh google credentials.", e);
+      return new HeaderCacheElement(
+          Status.UNAUTHENTICATED
+              .withDescription("TimeoutException during Authentication.")
+              .withCause(e)
+      );
+    } catch (Exception e) {
+      LOG.warn("Unexpected execption while trying to refresh google credentials.", e);
+      return new HeaderCacheElement(
+          Status.UNAUTHENTICATED
+              .withDescription("Unexpected execption during Authentication.")
+              .withCause(e)
+      );
+    }
+  }
+
+  /**
+   * Refreshes the OAuth2 token asynchronously. This method will only start an async refresh if
+   * there isn't a currently running asynchronous refresh and the current token is not "Good".
+   */
+  Future<HeaderCacheElement> asyncRefresh() {
+    LOG.trace("asyncRefresh");
+
+    synchronized (lock) {
+      try {
+        if (futureToken != null) {
+          return futureToken;
+        }
+        if (headerCache.getCacheState() == CacheState.Good) {
+          return Futures.immediateFuture(headerCache);
+        }
+
+        Future<HeaderCacheElement> future = executor.submit(new Callable<HeaderCacheElement>() {
+          @Override
+          public HeaderCacheElement call() throws Exception {
+            return updateToken();
+          }
+        });
+
+        if (!future.isDone()) {
+          this.futureToken = future;
+        }
+        return future;
+      } catch (RuntimeException e) {
+        futureToken = null;
+        LOG.warn("Got an unexpected exception while trying to refresh google credentials.", e);
+        return Futures.immediateFuture(new HeaderCacheElement(
+          Status.UNAUTHENTICATED
+              .withDescription("Unexpected error trying to authenticate")
+              .withCause(e)));
+      }
+    }
+  }
+
+    /**
+     * Create a new token via {@link OAuth2Credentials#refreshAccessToken()}, and then update the cache.
+     * @return the new token
+     */
+  private HeaderCacheElement updateToken() {
+    HeaderCacheElement newToken;
+    try {
+      LOG.info("Refreshing the OAuth token");
+      newToken = new HeaderCacheElement(credentials.refreshAccessToken());
+    } catch (Exception e) {
+      LOG.warn("Got an unexpected exception while trying to refresh google credentials.", e);
+      newToken = new HeaderCacheElement(
+          Status.UNAUTHENTICATED
+              .withDescription("Unexpected error trying to authenticate")
+              .withCause(e)
+      );
+    }
+    synchronized (lock) {
+      // Update the token only if the new token is good or the old token is bad
+      if (newToken.isValid() || !headerCache.isValid()) {
+        headerCache = newToken;
+      } else {
+        LOG.warn("Failed to refresh the access token. Falling back to existing token. "
+            + "New token state: {}, status: {}", newToken.getCacheState(), newToken.getToken().status);
+      }
+      futureToken = null;
+
+      return headerCache;
+    }
+  }
+
+    /**
+     * Clear the cache.
+     *
+     * @param oldToken for a comparison.  Only revoke the cache if the oldToken matches the current one.
+     */
+  void revokeUnauthToken(HeaderToken oldToken) {
+    synchronized (lock) {
+      if (headerCache.getToken() == oldToken) {
+        LOG.warn("Got unauthenticated response from server, revoking the current token");
+        headerCache = EMPTY_HEADER;
+      } else {
+        LOG.info("Skipping revoke, since the revoked token has already changed");
+      }
+    }
+  }
+
+  @VisibleForTesting
+  boolean isRefreshing() {
+    synchronized(lock) {
+      return futureToken != null;
+    }
+  }
+}
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 27f20e321..699b6b5a0 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
@@ -15,13 +15,10 @@
  */
 package com.google.cloud.bigtable.grpc.io;
 
-import com.google.api.client.util.Clock;
-import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.OAuth2Credentials;
 import com.google.cloud.bigtable.config.Logger;
+import com.google.cloud.bigtable.grpc.io.OAuthCredentialsCache.HeaderToken;
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Preconditions;
-import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.RateLimiter;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
@@ -30,19 +27,11 @@ import io.grpc.ClientCall.Listener;
 import io.grpc.ClientInterceptor;
 import io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
 import io.grpc.ForwardingClientCallListener.SimpleForwardingClientCallListener;
-import io.opencensus.trace.Tracing;
 import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
 import io.grpc.Status;
 
-import java.io.Closeable;
-import java.util.concurrent.Callable;
-import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Future;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
-import javax.annotation.concurrent.GuardedBy;
 
 /**
  * Client interceptor that authenticates all calls by binding header data provided by a credential.
@@ -62,140 +51,39 @@ import javax.annotation.concurrent.GuardedBy;
  * @version $Id: $Id
  */
 public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor {
+  private static final Logger LOG = new Logger(RefreshingOAuth2CredentialsInterceptor.class);
 
   public static int TIMEOUT_SECONDS = 15;
 
-  private static final Logger LOG = new Logger(RefreshingOAuth2CredentialsInterceptor.class);
-  private static final HeaderCacheElement EMPTY_HEADER = new HeaderCacheElement(null, 0);
-
   @VisibleForTesting
   static final Metadata.Key<String> AUTHORIZATION_HEADER_KEY = Metadata.Key.of(
       "Authorization", Metadata.ASCII_STRING_MARSHALLER);
 
-  @VisibleForTesting
-  static Clock clock = Clock.SYSTEM;
-
-  /**
-   * <p>
-   * This enum describes the states of the OAuth header.
-   * </p>
-   * <ol>
-   * <li>Good - fine to use, and does not need to be refreshed.
-   * <li>Stale - fine to use, but requires an async refresh
-   * <li>Expired - Cannot be used. Wait for a new token to be loaded
-   * </ol>
-   */
-  @VisibleForTesting
-  enum CacheState {
-    Good(true), Stale(true),
-    Expired(false), Exception(false);
-
-    private boolean isValid;
-
-    CacheState(boolean isValid) {
-      this.isValid = isValid;
-    }
-
-    public boolean isValid() {
-      return isValid;
-    }
-  }
-
-  @VisibleForTesting
-  static class HeaderCacheElement {
-
-    /**
-     * This specifies how far in advance of a header expiration do we consider the token stale. The
-     * Stale state indicates that the interceptor needs to do an asynchronous refresh.
-     */
-    static final long TOKEN_STALENESS_MS = TimeUnit.MINUTES.toMillis(6);
-
-    /**
-     * After the token is "expired," the interceptor blocks gRPC calls. The Expired state indicates
-     * that the interceptor needs to do a synchronous refresh.
-     */
-    // 5 minutes as per https://github.com/google/google-auth-library-java/pull/95
-    static final long TOKEN_EXPIRES_MS = TimeUnit.MINUTES.toMillis(5);
-
-    final Status status;
-    final String header;
-    final long actualExpirationTimeMs;
-
-    private HeaderCacheElement(AccessToken token) {
-      this.status = Status.OK;
-      if (token.getExpirationTime() == null) {
-        actualExpirationTimeMs = Long.MAX_VALUE;
-      } else {
-        actualExpirationTimeMs = token.getExpirationTime().getTime();
-      }
-      header = "Bearer " + token.getTokenValue();
-    }
-
-    private HeaderCacheElement(String header, long actualExpirationTimeMs) {
-      this.status = Status.OK;
-      this.header = header;
-      this.actualExpirationTimeMs = actualExpirationTimeMs;
-    }
-
-    private HeaderCacheElement(Status errorStatus) {
-      Preconditions.checkArgument(!errorStatus.isOk(), "Error status can't be OK");
-      this.status = errorStatus;
-      this.header = null;
-      this.actualExpirationTimeMs = 0;
-    }
-
-    CacheState getCacheState() {
-      long now = clock.currentTimeMillis();
-
-      if (!status.isOk()) {
-        return CacheState.Exception;
-      } else if (actualExpirationTimeMs - TOKEN_EXPIRES_MS <= now) {
-        return CacheState.Expired;
-      } else if (actualExpirationTimeMs - TOKEN_STALENESS_MS <= now) {
-        return CacheState.Stale;
-      } else {
-        return CacheState.Good;
-      }
-    }
-
-    private boolean isValid() {
-      return getCacheState().isValid();
-    }
-  }
+  private final OAuthCredentialsCache store;
+  private final RateLimiter rateLimiter;
 
-  private static class UnAuthResponseListener<RespT> extends SimpleForwardingClientCallListener<RespT> {
+  private class UnAuthResponseListener<RespT> extends SimpleForwardingClientCallListener<RespT> {
 
-    private final HeaderCacheElement origToken;
-    private final RefreshingOAuth2CredentialsInterceptor interceptor;
+    private final HeaderToken origToken;
 
-    private UnAuthResponseListener(RefreshingOAuth2CredentialsInterceptor interceptor, Listener<RespT> delegate,
-                                   HeaderCacheElement origToken) {
+    private UnAuthResponseListener(Listener<RespT> delegate, HeaderToken origToken) {
       super(delegate);
       this.origToken = origToken;
-      this.interceptor = interceptor;
     }
 
     @Override
     public void onClose(Status status, Metadata trailers) {
       if (status == Status.UNAUTHENTICATED) {
-        interceptor.revokeUnauthToken(origToken);
+        if (rateLimiter.tryAcquire()) {
+          store.revokeUnauthToken(origToken);
+        } else {
+          LOG.trace("UnAuthResponseListener rate limited");
+        }
       }
       super.onClose(status, trailers);
     }
   }
 
-  private final ExecutorService executor;
-  private final RateLimiter rateLimiter;
-  private final OAuth2Credentials credentials;
-  private 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
-  @GuardedBy("lock")
-  private volatile HeaderCacheElement headerCache = EMPTY_HEADER;
-
   /**
    * <p>Constructor for RefreshingOAuth2CredentialsInterceptor.</p>
    *
@@ -203,10 +91,14 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
    * @param credentials a {@link OAuth2Credentials} object.
    */
   public RefreshingOAuth2CredentialsInterceptor(ExecutorService scheduler, OAuth2Credentials credentials) {
-    this.executor = Preconditions.checkNotNull(scheduler);
-    this.credentials = Preconditions.checkNotNull(credentials);
+    this(new OAuthCredentialsCache(scheduler, credentials));
+  }
+
+  @VisibleForTesting
+  RefreshingOAuth2CredentialsInterceptor(OAuthCredentialsCache store) {
+    this.store = store;
 
-    // Retry every 10 seconds.
+    // Revoke credentials no more than every 10 seconds fo UNAUTHENTICATED statuses
     rateLimiter = RateLimiter.create(0.1);
   }
 
@@ -227,18 +119,17 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
 
       @Override
       public void start(Listener<RespT> responseListener, Metadata headers) {
-        HeaderCacheElement headerCache = getHeaderSafe();
+        HeaderToken token = store.getHeader(TIMEOUT_SECONDS);
 
-        if (!headerCache.status.isOk()) {
+        if (!token.getStatus().isOk()) {
           unauthorized = true;
-          responseListener.onClose(headerCache.status, new Metadata());
+          responseListener.onClose(token.getStatus(), new Metadata());
           return;
         }
 
-        headers.put(AUTHORIZATION_HEADER_KEY, headerCache.header);
+        headers.put(AUTHORIZATION_HEADER_KEY, token.getHeader());
 
-        delegate().start(new UnAuthResponseListener<>(RefreshingOAuth2CredentialsInterceptor.this,
-            responseListener, headerCache), headers);
+        delegate().start(new UnAuthResponseListener<>(responseListener, token), headers);
       }
 
       @Override
@@ -271,177 +162,11 @@ public class RefreshingOAuth2CredentialsInterceptor implements ClientInterceptor
     };
   }
 
-  @VisibleForTesting
-  CacheState getCacheState() {
-    return headerCache.getCacheState();
-  }
-
-  @VisibleForTesting
-  HeaderCacheElement getHeaderSafe() {
-    try {
-      return getHeader();
-    } catch (Exception e) {
-      LOG.warn("Got an unexpected exception while trying to refresh google credentials.", e);
-      return new HeaderCacheElement(
-          Status.UNAUTHENTICATED
-              .withDescription("Unexpected failure get auth token")
-              .withCause(e)
-      );
-    }
-  }
-
-  /**
-   * Get the http credential header we need from a new oauth2 AccessToken.
-   */
-  private HeaderCacheElement getHeader() {
-    // Optimize for the common case: do a volatile read to peek for a Good cache value
-    HeaderCacheElement headerCacheUnsync = this.headerCache;
-
-    // TODO(igorbernstein2): figure out how to make this work with appengine request scoped threads
-    switch (headerCacheUnsync.getCacheState()) {
-      case Good:
-        return headerCacheUnsync;
-      case Stale:
-        asyncRefresh();
-        return headerCacheUnsync;
-      case Expired:
-      case Exception:
-        // defer the future resolution (asyncRefresh will spin up a thread that will try to acquire the lock)
-        return syncRefresh();
-      default:
-        String message = "Could not process state: " + headerCacheUnsync.getCacheState();
-        LOG.warn(message);
-        return new HeaderCacheElement(
-            Status.UNAUTHENTICATED
-                .withCause(new IllegalStateException(message)));
-    }
-  }
-
-  /**
-   * Refresh the credentials and block. Will return an error if the credentials haven't
-   * been refreshed.
-   * This method should not be called while holding the refresh lock
-   */
-  private HeaderCacheElement syncRefresh() {
-    try (Closeable ss = Tracing.getTracer().spanBuilder("CredentialsRefresh").startScopedSpan()) {
-      return asyncRefresh().get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
-    } catch (InterruptedException e) {
-      LOG.warn("Interrupted while trying to refresh google credentials.", e);
-      Thread.currentThread().interrupt();
-      return new HeaderCacheElement(
-          Status.UNAUTHENTICATED
-              .withDescription("Authentication was interrupted.")
-              .withCause(e)
-      );
-    } catch (ExecutionException e) {
-      LOG.warn("ExecutionException while trying to refresh google credentials.", e);
-      return new HeaderCacheElement(
-          Status.UNAUTHENTICATED
-              .withDescription("ExecutionException during Authentication.")
-              .withCause(e)
-      );
-    } catch (TimeoutException e) {
-      LOG.warn("TimeoutException while trying to refresh google credentials.", e);
-      return new HeaderCacheElement(
-          Status.UNAUTHENTICATED
-              .withDescription("TimeoutException during Authentication.")
-              .withCause(e)
-      );
-    } catch (Exception e) {
-      LOG.warn("Unexpected execption while trying to refresh google credentials.", e);
-      return new HeaderCacheElement(
-          Status.UNAUTHENTICATED
-              .withDescription("Unexpected execption during Authentication.")
-              .withCause(e)
-      );
-    }
-  }
-
   /**
    * Refreshes the OAuth2 token asynchronously. This method will only start an async refresh if
    * there isn't a currently running asynchronous refresh.
    */
-  Future<HeaderCacheElement> asyncRefresh() {
-    LOG.trace("asyncRefresh");
-
-    synchronized (lock) {
-      try {
-        if (futureToken != null) {
-          return futureToken;
-        }
-        if (headerCache.getCacheState() == CacheState.Good) {
-          return Futures.immediateFuture(headerCache);
-        }
-
-        Future<HeaderCacheElement> future = executor.submit(new Callable<HeaderCacheElement>() {
-          @Override
-          public HeaderCacheElement call() throws Exception {
-            return updateToken();
-          }
-        });
-
-        if (!future.isDone()) {
-          this.futureToken = future;
-        }
-        return future;
-      } catch (RuntimeException e) {
-        futureToken = null;
-        LOG.warn("Got an unexpected exception while trying to refresh google credentials.", e);
-        return Futures.immediateFuture(new HeaderCacheElement(
-          Status.UNAUTHENTICATED
-              .withDescription("Unexpected error trying to authenticate")
-              .withCause(e)));
-      }
-    }
-  }
-
-  private HeaderCacheElement updateToken() {
-    HeaderCacheElement newToken;
-    try {
-      LOG.info("Refreshing the OAuth token");
-      newToken = new HeaderCacheElement(credentials.refreshAccessToken());
-    } catch (Exception e) {
-      LOG.warn("Got an unexpected exception while trying to refresh google credentials.", e);
-      newToken = new HeaderCacheElement(
-          Status.UNAUTHENTICATED
-              .withDescription("Unexpected error trying to authenticate")
-              .withCause(e)
-      );
-    }
-    synchronized (lock) {
-      // Update the token only if the new token is good or the old token is bad
-      if (newToken.isValid() || !headerCache.isValid()) {
-        headerCache = newToken;
-      } else {
-        LOG.warn("Failed to refresh the access token. Falling back to existing token. "
-            + "New token state: {}, status: {}", newToken.getCacheState(), newToken.status);
-      }
-      futureToken = null;
-
-      return headerCache;
-    }
-  }
-
-  private void revokeUnauthToken(HeaderCacheElement oldToken) {
-    if (!rateLimiter.tryAcquire()) {
-      LOG.trace("Rate limited");
-      return;
-    }
-
-    synchronized (lock) {
-      if (headerCache == oldToken) {
-        LOG.warn("Got unauthenticated response from server, revoking the current token");
-        headerCache = EMPTY_HEADER;
-      } else {
-        LOG.info("Skipping revoke, since the revoked token has already changed");
-      }
-    }
-  }
-
-  @VisibleForTesting
-  boolean isRefreshing() {
-    synchronized(lock) {
-      return futureToken != null;
-    }
+  void asyncRefresh() {
+    store.asyncRefresh();
   }
 }


Test patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/io/OAuthCredentialsCacheTest.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/io/OAuthCredentialsCacheTest.java
new file mode 100644
index 000000000..510522493
--- /dev/null
+++ b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/io/OAuthCredentialsCacheTest.java
@@ -0,0 +1,323 @@
+/*
+ * Copyright 2018 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable.grpc.io;
+
+import com.google.api.client.util.Clock;
+import com.google.auth.oauth2.AccessToken;
+import com.google.auth.oauth2.OAuth2Credentials;
+import com.google.cloud.bigtable.grpc.io.OAuthCredentialsCache.CacheState;
+import com.google.cloud.bigtable.grpc.io.OAuthCredentialsCache.HeaderCacheElement;
+import com.google.common.util.concurrent.SettableFuture;
+import io.grpc.*;
+import org.junit.*;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
+
+import java.io.IOException;
+import java.util.Date;
+import java.util.concurrent.*;
+
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.*;
+
+
+/**
+ * Tests for {@link OAuthCredentialsCache}
+ */
+@RunWith(JUnit4.class)
+public class OAuthCredentialsCacheTest {
+  private static int TIMEOUT_SECONDS = 5;
+
+  private static ExecutorService executorService;
+
+  @BeforeClass
+  public static void setup() {
+    executorService = Executors.newCachedThreadPool();
+  }
+
+  @AfterClass
+  public static void shtudown() {
+    executorService.shutdownNow();
+  }
+
+  private OAuthCredentialsCache underTest;
+
+  @Mock
+  private OAuth2Credentials mockCredentials;
+
+  @Mock
+  private Channel mockChannel;
+
+  @Mock
+  private ClientCall mockClientCall;
+
+  @Mock
+  private ClientCall.Listener mockListener;
+
+  @Before
+  public void setupMocks() {
+    MockitoAnnotations.initMocks(this);
+    setTimeInMillieconds(0L);
+    when(mockChannel.newCall(any(MethodDescriptor.class), any(CallOptions.class)))
+        .thenReturn(mockClientCall);
+  }
+
+  private void setTimeInMillieconds(final long timeMs) {
+    OAuthCredentialsCache.clock = new Clock() {
+      @Override
+      public long currentTimeMillis() {
+        return timeMs;
+      }
+    };
+  }
+
+  @Test
+  public void testSyncRefresh() throws Exception {
+    initialize(HeaderCacheElement.TOKEN_STALENESS_MS + 1);
+    Assert.assertEquals(CacheState.Good, underTest.getHeaderCache().getCacheState());
+  }
+
+  @SuppressWarnings("unchecked")
+  @Test
+  public void testExecutorProblem() throws Exception {
+    ExecutorService mockExecutor = Mockito.mock(ExecutorService.class);
+    when(mockExecutor.submit(any(Callable.class))).thenThrow(new RuntimeException(""));
+    underTest = new OAuthCredentialsCache(mockExecutor, mockCredentials);
+    Assert.assertFalse(underTest.getHeader(TIMEOUT_SECONDS).getStatus().isOk());
+  }
+
+  @Test
+  public void testStaleAndExpired() throws Exception {
+    long expiration = HeaderCacheElement.TOKEN_STALENESS_MS + 1;
+    initialize(expiration);
+    Assert.assertEquals(CacheState.Good, underTest.getHeaderCache().getCacheState());
+    long startTime = 2L;
+    setTimeInMillieconds(startTime);
+    Assert.assertEquals(CacheState.Stale, underTest.getHeaderCache().getCacheState());
+    long expiredStaleDiff =
+        HeaderCacheElement.TOKEN_STALENESS_MS - HeaderCacheElement.TOKEN_EXPIRES_MS;
+    setTimeInMillieconds(startTime + expiredStaleDiff);
+    Assert.assertEquals(CacheState.Expired, underTest.getHeaderCache().getCacheState());
+  }
+
+  @Test
+  public void testNullExpiration() throws Exception {
+    setTimeInMillieconds(100);
+    Mockito.when(mockCredentials.refreshAccessToken()).thenReturn(
+        new AccessToken("", null));
+    underTest = new OAuthCredentialsCache(executorService, mockCredentials);
+    underTest.asyncRefresh().get(100, TimeUnit.MILLISECONDS);
+    Assert.assertEquals(CacheState.Good, underTest.getHeaderCache().getCacheState());
+
+    // Make sure that the getHeader() doesn't expire in the distant future.
+    setTimeInMillieconds(100000000000L);
+    Assert.assertEquals(CacheState.Good, underTest.getHeaderCache().getCacheState());
+  }
+
+  @Test
+  public void testRefreshAfterFailure() throws Exception {
+    underTest = new OAuthCredentialsCache(executorService, mockCredentials);
+
+    final AccessToken accessToken = new AccessToken("hi", new Date(HeaderCacheElement.TOKEN_STALENESS_MS + 1));
+
+    //noinspection unchecked
+    Mockito.when(mockCredentials.refreshAccessToken())
+        // First call will throw Exception & bypass retries
+        .thenThrow(new IOException())
+        // Second call will succeed
+        .thenReturn(accessToken);
+
+    // First call
+    OAuthCredentialsCache.HeaderToken token1 = underTest.getHeader(TIMEOUT_SECONDS);
+    Assert.assertFalse(token1.getStatus().isOk());
+
+    // Now the second token should be available
+    OAuthCredentialsCache.HeaderToken token2 = underTest.getHeader(TIMEOUT_SECONDS);
+    Assert.assertTrue(token2.getStatus().isOk());
+    Assert.assertThat(token2.getHeader(), containsString("hi"));
+    // Make sure that the token was only requested twice: once for the first failure & second time for background recovery
+    Mockito.verify(mockCredentials, times(2)).refreshAccessToken();
+  }
+
+  @Test
+  public void testRefreshAfterStale() throws Exception {
+    underTest = new OAuthCredentialsCache(executorService, mockCredentials);
+
+    long expires = HeaderCacheElement.TOKEN_STALENESS_MS + 1;
+    final AccessToken staleToken = new AccessToken("stale", new Date(expires));
+    AccessToken goodToken =
+        new AccessToken("good", new Date(expires + HeaderCacheElement.TOKEN_STALENESS_MS + 1));
+
+    //noinspection unchecked
+    Mockito.when(mockCredentials.refreshAccessToken())
+        // First call will setup a stale token
+        .thenReturn(staleToken)
+        // Second call will give a good token
+        .thenReturn(goodToken);
+
+    // First call - setup
+    OAuthCredentialsCache.HeaderToken firstResult = underTest.getHeader(TIMEOUT_SECONDS);
+    Assert.assertEquals(CacheState.Good, underTest.getHeaderCache().getCacheState());
+    Assert.assertThat(firstResult.getHeader(), containsString("stale"));
+
+    // Fast forward until token is stale
+    setTimeInMillieconds(10);
+
+    // Second call - return stale token, but schedule refresh
+    OAuthCredentialsCache.HeaderToken secondResult = underTest.getHeader(TIMEOUT_SECONDS);
+    Assert.assertEquals(CacheState.Stale, underTest.getHeaderCache().getCacheState());
+    Assert.assertThat(secondResult.getHeader(), containsString("stale"));
+
+    // forward to expired
+    setTimeInMillieconds(expires);
+
+    // Third call - now returns good token
+    OAuthCredentialsCache.HeaderToken thirdResult = underTest.getHeader(TIMEOUT_SECONDS);
+    Assert.assertEquals(CacheState.Good, underTest.getHeaderCache().getCacheState());
+    Assert.assertThat(thirdResult.getHeader(), containsString("good"));
+
+    // Make sure that the token was only requested twice: once for the stale token & second time for the good token
+    Mockito.verify(mockCredentials, times(2)).refreshAccessToken();
+  }
+
+  @Test(timeout = 30000)
+  /*
+   * Test that checks that concurrent requests to RefreshingOAuth2CredentialsInterceptor refresh
+   * logic doesn't cause hanging behavior.  Specifically, when an Expired condition occurs it
+   * triggers a call to syncRefresh() which potentially waits for refresh that was initiated
+   * from another thread either through syncRefresh() or asyncRefresh().  This test case simulates
+   * that condition.
+   */
+  public void testRefreshDoesntHang() throws Exception {
+    for (int i = 0; i < 100; i++) {
+      MockitoAnnotations.initMocks(this);
+      testHanging();
+    }
+  }
+
+  @Test
+  public void testRevoke() throws Exception {
+    long expiration = HeaderCacheElement.TOKEN_STALENESS_MS + 1;
+    initialize(expiration);
+    Assert.assertEquals(CacheState.Good, underTest.getHeaderCache().getCacheState());
+
+    underTest.revokeUnauthToken(underTest.getHeaderCache().getToken());
+    Assert.assertEquals(CacheState.Expired, underTest.getHeaderCache().getCacheState());
+
+    underTest.getHeader(TIMEOUT_SECONDS);
+    Assert.assertEquals(CacheState.Good, underTest.getHeaderCache().getCacheState());
+
+    underTest.revokeUnauthToken(underTest.getHeaderCache().getToken());
+    // Ensure that the revoke interceptor does not run so quickly after the previous revoke.
+    // The cache state should stay good.
+    Assert.assertTrue(underTest.getHeader(5).getStatus().isOk());
+  }
+
+  private void testHanging()
+      throws Exception, InterruptedException, ExecutionException, TimeoutException {
+    // Assume that the user starts at this time... it's an arbitrarily big number which will
+    // assure that subtracting HeaderCacheElement.TOKEN_STALENESS_MS and TOKEN_EXPIRES_MS will not
+    // be negative.
+    long start = HeaderCacheElement.TOKEN_STALENESS_MS * 10;
+    setTimeInMillieconds(start);
+
+    // RefreshingOAuth2CredentialsInterceptor will show that the access token is stale.
+    final long expiration = start + HeaderCacheElement.TOKEN_EXPIRES_MS + 1;
+
+    // Create a mechanism that will allow us to control when the accessToken is returned.
+    // mockCredentials.refreshAccessToken() will get called asynchronously and will wait until the
+    // lock is notified before returning.  That will allow us to set up multiple concurrent calls
+    final FutureAnswer<AccessToken> answer = new FutureAnswer<>();
+    Mockito.when(mockCredentials.refreshAccessToken()).thenAnswer(answer);
+
+    underTest =
+        new OAuthCredentialsCache(executorService, mockCredentials);
+
+    // At this point, the access token wasn't retrieved yet. The
+    // RefreshingOAuth2CredentialsInterceptor considers null to be Expired.
+    Assert.assertEquals(CacheState.Expired, underTest.getHeaderCache().getCacheState());
+
+    syncCall(answer, new Date(expiration));
+
+    // Check to make sure that the AccessToken was retrieved.
+    Assert.assertEquals(CacheState.Stale, underTest.getHeaderCache().getCacheState());
+
+    // Check to make sure we're no longer refreshing.
+    Assert.assertFalse(underTest.isRefreshing());
+
+    answer.reset();
+
+    // Kick off 100 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
+    // requests that call asyncRefresh() when the token turns stale.
+    Future<HeaderCacheElement> previous = underTest.asyncRefresh();
+    for (int i = 0; i < 100; i++) {
+      Future<HeaderCacheElement> current = underTest.asyncRefresh();
+      Assert.assertEquals(previous, current);
+      previous = current;
+    }
+
+    syncCall(answer, new Date(expiration + HeaderCacheElement.TOKEN_EXPIRES_MS + 1));
+    Assert.assertFalse(underTest.isRefreshing());
+  }
+
+  private static class FutureAnswer<T> implements Answer<T> {
+
+    private SettableFuture<T> future = SettableFuture.create();
+
+    @Override
+    public T answer(InvocationOnMock invocation) throws Throwable {
+      return future.get();
+    }
+
+    void set(T value) {
+      future.set(value);
+    }
+
+    void reset() {
+      future = SettableFuture.create();
+    }
+  }
+
+  private void syncCall(FutureAnswer<AccessToken> answer, Date expirationTime)
+      throws InterruptedException, ExecutionException, TimeoutException {
+    underTest.asyncRefresh();
+
+    // There should be a single thread kicked off by the underTest.asyncRefresh() calls about
+    // 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 refresh().. Make sure
+    // that at least 1 refresh process is in progress.
+    Assert.assertTrue(underTest.isRefreshing());
+
+    answer.set(new AccessToken("hi", expirationTime));
+    underTest.asyncRefresh().get(100, TimeUnit.MILLISECONDS);
+  }
+
+  private void initialize(long expiration) throws Exception, ExecutionException, InterruptedException {
+    Mockito.when(mockCredentials.refreshAccessToken()).thenReturn(
+        new AccessToken("", new Date(expiration)));
+    underTest = new OAuthCredentialsCache(executorService, mockCredentials);
+    Assert.assertEquals(CacheState.Expired, underTest.getHeaderCache().getCacheState());
+    underTest.asyncRefresh().get();
+    Assert.assertFalse(underTest.isRefreshing());
+  }
+}
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 d4a29c820..92b0a52df 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,34 +15,17 @@
  */
 package com.google.cloud.bigtable.grpc.io;
 
-import static org.hamcrest.CoreMatchers.containsString;
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.eq;
-import static org.mockito.Matchers.same;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
-import com.google.api.client.util.Clock;
+import static org.mockito.Matchers.*;
+import static org.mockito.Mockito.*;
+
 import com.google.bigtable.v2.BigtableGrpc;
 import com.google.bigtable.v2.ReadRowsRequest;
 import com.google.bigtable.v2.ReadRowsResponse;
-import com.google.common.util.concurrent.MoreExecutors;
-import com.google.common.util.concurrent.SettableFuture;
-import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.OAuth2Credentials;
-import com.google.cloud.bigtable.grpc.io.RefreshingOAuth2CredentialsInterceptor.CacheState;
-import com.google.cloud.bigtable.grpc.io.RefreshingOAuth2CredentialsInterceptor.HeaderCacheElement;
 
 import java.io.IOException;
-import java.util.Date;
-import java.util.concurrent.Callable;
-import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 
 import io.grpc.*;
 import org.junit.AfterClass;
@@ -54,10 +37,7 @@ import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
-import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
-import org.mockito.invocation.InvocationOnMock;
-import org.mockito.stubbing.Answer;
 
 
 /**
@@ -66,320 +46,155 @@ import org.mockito.stubbing.Answer;
 @RunWith(JUnit4.class)
 public class RefreshingOAuth2CredentialsInterceptorTest {
 
-  private static ExecutorService executorService;
-
-  @BeforeClass
-  public static void setup() {
-    executorService = Executors.newCachedThreadPool();
-  }
-
-  @AfterClass
-  public static void shtudown() {
-    executorService.shutdownNow();
-  }
-
-  private RefreshingOAuth2CredentialsInterceptor underTest;
-
-  @Mock
-  private OAuth2Credentials mockCredentials;
-
-  @Mock
-  private Channel mockChannel;
-
-  @Mock
-  private ClientCall mockClientCall;
-
-  @Mock
-  private ClientCall.Listener mockListener;
-
-  @Before
-  public void setupMocks() {
-    MockitoAnnotations.initMocks(this);
-    setTimeInMillieconds(0L);
-    when(mockChannel.newCall(any(MethodDescriptor.class), any(CallOptions.class)))
-        .thenReturn(mockClientCall);
-  }
-
-  private void setTimeInMillieconds(final long timeMs) {
-    RefreshingOAuth2CredentialsInterceptor.clock = new Clock() {
-      @Override
-      public long currentTimeMillis() {
-        return timeMs;
-      }
-    };
-  }
-
-  @Test
-  public void testSyncRefresh() throws IOException {
-    Mockito.when(mockCredentials.refreshAccessToken()).thenReturn(
-        new AccessToken("", new Date(HeaderCacheElement.TOKEN_STALENESS_MS + 1)));
-    underTest = new RefreshingOAuth2CredentialsInterceptor(executorService, mockCredentials);
-    Assert.assertEquals(CacheState.Expired, underTest.getCacheState());
-    underTest.getHeaderSafe();
-    Assert.assertNotEquals(CacheState.Exception, underTest.getCacheState());
-    Assert.assertEquals(CacheState.Good, underTest.getCacheState());
-    Assert.assertFalse(underTest.isRefreshing());
-  }
-
-  @Test
-  /**
-   * Basic test to make sure that the interceptor works properly
-   */
-  public void testIntercept() throws IOException {
-    initialize(HeaderCacheElement.TOKEN_STALENESS_MS + 1);
-    initializeCall(CallOptions.DEFAULT);
-  }
-
-  @Test
-  /**
-   * Test to make sure that the interceptor revokes credentials if an RPC receives an UNAUTHENTICATED status.
-   */
-  public void testInvalid() throws IOException {
-    long expiration = HeaderCacheElement.TOKEN_STALENESS_MS + 1;
-    initialize(expiration);
-    ClientCall.Listener listener = initializeCall(CallOptions.DEFAULT);
-    listener.onClose(Status.UNAUTHENTICATED, new Metadata());
-    Assert.assertEquals(CacheState.Expired, underTest.getCacheState());
-  }
-
-  private ClientCall.Listener initializeCall(CallOptions callOptions) throws IOException {
-    ClientCall<ReadRowsRequest, ReadRowsResponse> call = underTest.interceptCall(
-        BigtableGrpc.METHOD_READ_ROWS,
-        callOptions,
-        mockChannel);
-    Metadata metadata = new Metadata();
-    call.start(mockListener, metadata);
-    Assert.assertEquals(underTest.getHeaderSafe().header,
-        metadata.get(RefreshingOAuth2CredentialsInterceptor.AUTHORIZATION_HEADER_KEY));
-
-    ArgumentCaptor<ClientCall.Listener> listenerCaptor = ArgumentCaptor.forClass(ClientCall.Listener.class);
-    verify(mockClientCall, times(1))
-        .start(listenerCaptor.capture(), same(metadata));
-
-    call.request(1);
-    call.sendMessage(ReadRowsRequest.getDefaultInstance());
-
-    verify(mockClientCall, times(1))
-        .request(eq(1));
-    verify(mockClientCall, times(1))
-        .sendMessage(same(ReadRowsRequest.getDefaultInstance()));
-
-    return listenerCaptor.getValue();
-  }
-
-  @Test
-  public void testRefresh() throws IOException {
-    Mockito.when(mockCredentials.refreshAccessToken()).thenReturn(
-        new AccessToken("", new Date(HeaderCacheElement.TOKEN_STALENESS_MS + 1)));
-    underTest = new RefreshingOAuth2CredentialsInterceptor(MoreExecutors.newDirectExecutorService(),
-        mockCredentials);
-    underTest.getHeaderSafe();
-    Assert.assertEquals(CacheState.Good, underTest.getCacheState());
-    Assert.assertFalse(underTest.isRefreshing());
-  }
-
-  @SuppressWarnings("unchecked")
-  @Test
-  public void testExecutorProblem() throws IOException {
-    ExecutorService mockExecutor = Mockito.mock(ExecutorService.class);
-    when(mockExecutor.submit(any(Callable.class))).thenThrow(new RuntimeException(""));
-    underTest = new RefreshingOAuth2CredentialsInterceptor(mockExecutor, mockCredentials);
-    Assert.assertEquals(CacheState.Exception, underTest.getHeaderSafe().getCacheState());
-  }
-
-  @Test
-  public void testStaleAndExpired() throws IOException {
-    long expiration = HeaderCacheElement.TOKEN_STALENESS_MS + 1;
-    initialize(expiration);
-    Assert.assertEquals(CacheState.Good, underTest.getCacheState());
-    long startTime = 2L;
-    setTimeInMillieconds(startTime);
-    Assert.assertEquals(CacheState.Stale, underTest.getCacheState());
-    long expiredStaleDiff =
-        HeaderCacheElement.TOKEN_STALENESS_MS - HeaderCacheElement.TOKEN_EXPIRES_MS;
-    setTimeInMillieconds(startTime + expiredStaleDiff);
-    Assert.assertEquals(CacheState.Expired, underTest.getCacheState());
-  }
-
-  @Test
-  public void testNullExpiration() throws Exception {
-    setTimeInMillieconds(100);
-    Mockito.when(mockCredentials.refreshAccessToken()).thenReturn(
-        new AccessToken("", null));
-    underTest = new RefreshingOAuth2CredentialsInterceptor(executorService, mockCredentials);
-    underTest.asyncRefresh().get(100, TimeUnit.MILLISECONDS);
-    Assert.assertEquals(CacheState.Good, underTest.getCacheState());
-
-    // Make sure that the header doesn't expire in the distant future.
-    setTimeInMillieconds(100000000000L);
-    Assert.assertEquals(CacheState.Good, underTest.getCacheState());
-  }
-
-  @Test
-  public void testRefreshAfterFailure() throws Exception {
-    underTest = new RefreshingOAuth2CredentialsInterceptor(executorService, mockCredentials);
-
-    final AccessToken accessToken = new AccessToken("hi", new Date(HeaderCacheElement.TOKEN_STALENESS_MS + 1));
-
-    //noinspection unchecked
-    Mockito.when(mockCredentials.refreshAccessToken())
-        // First call will throw Exception & bypass retries
-        .thenThrow(new IOException())
-        // Second call will succeed
-        .thenReturn(accessToken);
-
-    // First call
-    HeaderCacheElement firstResult = underTest.getHeaderSafe();
-    Assert.assertEquals(CacheState.Exception, firstResult.getCacheState());
-
-    // Now the second token should be available
-    HeaderCacheElement secondResult = underTest.getHeaderSafe();
-    Assert.assertEquals(CacheState.Good, secondResult.getCacheState());
-    Assert.assertThat(secondResult.header, containsString("hi"));
-    // Make sure that the token was only requested twice: once for the first failure & second time for background recovery
-    Mockito.verify(mockCredentials, times(2)).refreshAccessToken();
-  }
-
-  @Test
-  public void testRefreshAfterStale() throws Exception {
-    underTest = new RefreshingOAuth2CredentialsInterceptor(executorService, mockCredentials);
-
-    long expires = HeaderCacheElement.TOKEN_STALENESS_MS + 1;
-    final AccessToken staleToken = new AccessToken("stale", new Date(expires));
-    AccessToken goodToken =
-        new AccessToken("good", new Date(expires + HeaderCacheElement.TOKEN_STALENESS_MS + 1));
-
-    //noinspection unchecked
-    Mockito.when(mockCredentials.refreshAccessToken())
-        // First call will setup a stale token
-        .thenReturn(staleToken)
-        // Second call will give a good token
-        .thenReturn(goodToken);
-
-    // First call - setup
-    HeaderCacheElement firstResult = underTest.getHeaderSafe();
-    Assert.assertEquals(CacheState.Good, firstResult.getCacheState());
-    Assert.assertThat(firstResult.header, containsString("stale"));
-
-    // Fast forward until token is stale
-    setTimeInMillieconds(10);
-
-    // Second call - return stale token, but schedule refresh
-    HeaderCacheElement secondResult = underTest.getHeaderSafe();
-    Assert.assertEquals(CacheState.Stale, secondResult.getCacheState());
-    Assert.assertThat(secondResult.header, containsString("stale"));
-
-    // forward to expired
-    setTimeInMillieconds(expires);
-
-    // Third call - now returns good token
-    HeaderCacheElement thirdResult = underTest.getHeaderSafe();
-    Assert.assertEquals(CacheState.Good, thirdResult.getCacheState());
-    Assert.assertThat(thirdResult.header, containsString("good"));
-
-    // Make sure that the token was only requested twice: once for the stale token & second time for the good token
-    Mockito.verify(mockCredentials, times(2)).refreshAccessToken();
-  }
-
-  @Test(timeout = 30000)
-  /*
-   * Test that checks that concurrent requests to RefreshingOAuth2CredentialsInterceptor refresh
-   * logic doesn't cause hanging behavior.  Specifically, when an Expired condition occurs it
-   * triggers a call to syncRefresh() which potentially waits for refresh that was initiated
-   * from another thread either through syncRefresh() or asyncRefresh().  This test case simulates
-   * that condition.
-   */
-  public void testRefreshDoesntHang() throws Exception {
-    for (int i = 0; i < 100; i++) {
-      MockitoAnnotations.initMocks(this);
-      testHanging();
+    private static ExecutorService executorService;
+    private static String HEADER = "SomeHeader";
+
+    @BeforeClass
+    public static void setup() {
+        executorService = Executors.newCachedThreadPool();
+    }
+
+    @AfterClass
+    public static void shtudown() {
+        executorService.shutdownNow();
+    }
+
+    private RefreshingOAuth2CredentialsInterceptor underTest;
+
+    @Mock
+    private OAuth2Credentials mockCredentials;
+
+    @Mock
+    private Channel mockChannel;
+
+    @Mock
+    private ClientCall mockClientCall;
+
+    @Mock
+    private ClientCall.Listener mockListener;
+
+    @Mock
+    private OAuthCredentialsCache mockStore;
+
+    @Before
+    public void setupMocks() {
+        MockitoAnnotations.initMocks(this);
+        when(mockChannel.newCall(any(MethodDescriptor.class), any(CallOptions.class)))
+                .thenReturn(mockClientCall);
+        underTest = new RefreshingOAuth2CredentialsInterceptor(mockStore);
+    }
+
+    @Test
+    /**
+     * Basic test to make sure that the interceptor works properly
+     */
+    public void testValidAuthToken() throws IOException {
+        initializeOk();
+        sendRequest(CallOptions.DEFAULT);
+        receiveMessage();
+        checkOKCompletedCorrectly();
     }
-  }
-
-  private void testHanging()
-      throws IOException, InterruptedException, ExecutionException, TimeoutException {
-    // Assume that the user starts at this time... it's an arbitrarily big number which will
-    // assure that subtracting HeaderCacheElement.TOKEN_STALENESS_MS and TOKEN_EXPIRES_MS will not
-    // be negative.
-    long start = HeaderCacheElement.TOKEN_STALENESS_MS * 10;
-    setTimeInMillieconds(start);
-
-    // RefreshingOAuth2CredentialsInterceptor will show that the access token is stale.
-    final long expiration = start + HeaderCacheElement.TOKEN_EXPIRES_MS + 1;
-
-    // Create a mechanism that will allow us to control when the accessToken is returned.
-    // mockCredentials.refreshAccessToken() will get called asynchronously and will wait until the
-    // lock is notified before returning.  That will allow us to set up multiple concurrent calls
-    final FutureAnswer<AccessToken> answer = new FutureAnswer<>();
-    Mockito.when(mockCredentials.refreshAccessToken()).thenAnswer(answer);
-
-    underTest =
-        new RefreshingOAuth2CredentialsInterceptor(executorService, mockCredentials);
-
-    // At this point, the access token wasn't retrieved yet. The
-    // RefreshingOAuth2CredentialsInterceptor considers null to be Expired.
-    Assert.assertEquals(CacheState.Expired, underTest.getCacheState());
-
-    syncCall(answer, new Date(expiration));
-
-    // Check to make sure that the AccessToken was retrieved.
-    Assert.assertEquals(CacheState.Stale, underTest.getCacheState());
-
-    // Check to make sure we're no longer refreshing.
-    Assert.assertFalse(underTest.isRefreshing());
-
-    answer.reset();
-
-    // Kick off 100 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
-    // requests that call asyncRefresh() when the token turns stale.
-    Future<HeaderCacheElement> previous = underTest.asyncRefresh();
-    for (int i = 0; i < 100; i++) {
-      Future<HeaderCacheElement> current = underTest.asyncRefresh();
-      Assert.assertEquals(previous, current);
-      previous = current;
+
+    @Test
+    /**
+     * Test to make sure that the interceptor revokes credentials if an RPC receives an UNAUTHENTICATED status.
+     */
+    public void testInvalidAuthToken() throws IOException {
+        Status grpcStatus = Status.UNAUTHENTICATED;
+
+        // Something bad happened, and authentication could not be established
+        when(mockStore.getHeader(anyInt()))
+                .thenReturn(new OAuthCredentialsCache.HeaderToken(grpcStatus, null));
+
+        ClientCall<ReadRowsRequest, ReadRowsResponse> call = underTest.interceptCall(
+                BigtableGrpc.METHOD_READ_ROWS,
+                CallOptions.DEFAULT,
+                mockChannel);
+        call.start(mockListener, new Metadata());
+
+        // Since the auth header is not available, the call should immediately fail with a non-OK status
+        verify(mockListener, times(1)).onClose(same(grpcStatus), any(Metadata.class));
+
+        // Users might do some additional work with call, so it shouldn't be null, but should not
+        // forward the request to the underlying call either.
+
+        // Verify that request, sendMessage, halfClose and cancel calls on the call object returned from
+        // is not null and does not
+        call.request(1);
+        verify(mockClientCall, times(0)).request(anyInt());
+
+        call.halfClose();
+        verify(mockClientCall, times(0)).halfClose();
+
+        call.sendMessage(ReadRowsRequest.getDefaultInstance());
+        verify(mockClientCall, times(0)).sendMessage(any(ReadRowsRequest.class));
+
+        call.cancel("", null);
+        verify(mockClientCall, times(0)).cancel(anyString(), any(Throwable.class));
     }
 
-    syncCall(answer, new Date(expiration + HeaderCacheElement.TOKEN_EXPIRES_MS + 1));
-    Assert.assertFalse(underTest.isRefreshing());
-  }
+    @Test
+    /**
+     * Test to make sure that the interceptor revokes credentials if an RPC receives an UNAUTHENTICATED status.
+     */
+    public void testRpcReturnsUnauthenticated() throws IOException {
+        initializeOk();
+
+        // 10 requests come in around the same time, and all of them get UNAUTHENTICATED
+        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(
+                any(OAuthCredentialsCache.HeaderToken.class));
+    }
 
-  private static class FutureAnswer<T> implements Answer<T> {
+    private ClientCall<ReadRowsRequest, ReadRowsResponse> sendRequest(CallOptions callOptions) {
+        ClientCall<ReadRowsRequest, ReadRowsResponse> call = underTest.interceptCall(
+                BigtableGrpc.METHOD_READ_ROWS,
+                callOptions,
+                mockChannel);
+        call.start(mockListener, new Metadata());
+        call.request(1);
+        call.sendMessage(ReadRowsRequest.getDefaultInstance());
+        call.halfClose();
+        return call;
+    }
+
+    private ClientCall.Listener receiveMessage() {
+        ClientCall.Listener listener = getListener();
+        listener.onMessage(ReadRowsResponse.getDefaultInstance());
+        listener.onClose(Status.OK, new Metadata());
+
+        return listener;
+    }
 
-    private SettableFuture<T> future = SettableFuture.create();
+    private ClientCall.Listener getListener() {
+        ArgumentCaptor<ClientCall.Listener> listenerCaptor = ArgumentCaptor.forClass(ClientCall.Listener.class);
+        verify(mockClientCall, atLeastOnce()).start(listenerCaptor.capture(), any(Metadata.class));
 
-    @Override
-    public T answer(InvocationOnMock invocation) throws Throwable {
-      return future.get();
+        // this gets the last matched value
+        return listenerCaptor.getValue();
     }
 
-    void set(T value) {
-      future.set(value);
+    private void initializeOk() {
+        when(mockStore.getHeader(anyInt()))
+                .thenReturn(new OAuthCredentialsCache.HeaderToken(Status.OK, HEADER));
     }
 
-    void reset() {
-      future = SettableFuture.create();
+    private void checkOKCompletedCorrectly() {
+        ArgumentCaptor<Metadata> metadataCaptor = ArgumentCaptor.forClass(Metadata.class);
+        verify(mockClientCall, atLeastOnce()).start(any(ClientCall.Listener.class), metadataCaptor.capture());
+        Assert.assertEquals(HEADER,
+                metadataCaptor.getValue().get(RefreshingOAuth2CredentialsInterceptor.AUTHORIZATION_HEADER_KEY));
+
+        verify(mockClientCall, times(1))
+                .request(eq(1));
+        verify(mockClientCall, times(1))
+                .sendMessage(same(ReadRowsRequest.getDefaultInstance()));
+        verify(mockListener, times(1)).onMessage(same(ReadRowsResponse.getDefaultInstance()));
+        verify(mockListener, times(1)).onClose(same(Status.OK), any(Metadata.class));
     }
-  }
-
-  private void syncCall(FutureAnswer<AccessToken> answer, Date expirationTime)
-      throws InterruptedException, ExecutionException, TimeoutException {
-    underTest.asyncRefresh();
-
-    // There should be a single thread kicked off by the underTest.asyncRefresh() calls about
-    // 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 refresh().. Make sure
-    // that at least 1 refresh process is in progress.
-    Assert.assertTrue(underTest.isRefreshing());
-
-    answer.set(new AccessToken("hi", expirationTime));
-    underTest.asyncRefresh().get(100, TimeUnit.MILLISECONDS);
-  }
-
-  private void initialize(long expiration) throws IOException {
-    Mockito.when(mockCredentials.refreshAccessToken()).thenReturn(
-        new AccessToken("", new Date(expiration)));
-    underTest = new RefreshingOAuth2CredentialsInterceptor(executorService, mockCredentials);
-    underTest.getHeaderSafe();
-  }
 }