OAuth2Client: improve sleep/close detection (#8337)

projectnessie · Apr 16, 2024 · b86d8de · b86d8de
1 parent 0d08bdd
commit b86d8de
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/api/client/src/main/java/org/projectnessie/client/auth/oauth2/OAuth2Client.java b/api/client/src/main/java/org/projectnessie/client/auth/oauth2/OAuth2Client.java
@@ -82,7 +82,6 @@ class OAuth2Client implements OAuth2Authenticator, Closeable {
             .getExecutor()
             .orElseGet(
                 () -> new OAuth2TokenRefreshExecutor(config.getBackgroundThreadIdleTimeout()));
-    lastAccess = config.getClock().get();
     currentTokensStage = started.thenApplyAsync((v) -> fetchNewTokens(), this.executor);
     currentTokensStage
         .whenComplete((tokens, error) -> log(error))
@@ -128,6 +127,7 @@ private Tokens getCurrentTokensIfAvailable() {
 
   @Override
   public void start() {
+    lastAccess = config.getClock().get();
     started.complete(null);
   }
 
@@ -158,6 +158,9 @@ public void close() {
   }
 
   private void wakeUp(Instant now) {
+    if (closing.get()) {
+      return;
+    }
     LOGGER.debug("Waking up...");
     Tokens currentTokens = getCurrentTokensIfAvailable();
     Duration delay = nextTokenRefresh(currentTokens, now, Duration.ZERO);
@@ -171,6 +174,9 @@ private void wakeUp(Instant now) {
   }
 
   private void maybeScheduleTokensRenewal(Tokens currentTokens) {
+    if (closing.get()) {
+      return;
+    }
     Instant now = config.getClock().get();
     if (Duration.between(lastAccess, now).compareTo(config.getPreemptiveTokenRefreshIdleTimeout())
         > 0) {
@@ -215,8 +221,7 @@ private void renewTokens() {
 
   private void log(Throwable error) {
     if (error != null) {
-      boolean tokensStageCancelled = error instanceof CancellationException && closing.get();
-      if (tokensStageCancelled) {
+      if (closing.get()) {
         return;
       }
       if (error instanceof CompletionException) {
@@ -250,10 +255,14 @@ Tokens fetchNewTokens() {
       case AUTHORIZATION_CODE:
         try (AuthorizationCodeFlow flow = new AuthorizationCodeFlow(config)) {
           return flow.fetchNewTokens();
+        } finally {
+          lastAccess = config.getClock().get();
         }
       case DEVICE_CODE:
         try (DeviceCodeFlow flow = new DeviceCodeFlow(config)) {
           return flow.fetchNewTokens();
+        } finally {
+          lastAccess = config.getClock().get();
         }
       default:
         throw new IllegalStateException("Unsupported grant type: " + config.getGrantType());