Use Instant for token expiration time

Author: acogoluegnes

src/main/java/com/rabbitmq/client/amqp/impl/TokenCredentials.java

@@ -32,6 +32,7 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Function;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -46,11 +47,13 @@ final class TokenCredentials implements Credentials {
   private final Map<Long, RegistrationImpl> registrations = new ConcurrentHashMap<>();
   private final AtomicLong registrationSequence = new AtomicLong(0);
   private final AtomicBoolean schedulingRenewal = new AtomicBoolean(false);
+  private final Function<Instant, Duration> refreshDelayStrategy;
   private volatile ScheduledFuture<?> renewalTask;
 
   TokenCredentials(TokenRequester requester, ScheduledExecutorService scheduledExecutorService) {
     this.requester = requester;
     this.scheduledExecutorService = scheduledExecutorService;
+    this.refreshDelayStrategy = new RatioRefreshDelayStrategy(0.8f);
   }
 
   private void lock() {
@@ -61,16 +64,8 @@ private void unlock() {
     this.lock.unlock();
   }
 
-  private boolean expiresSoon(Token t) {
-    // TODO use strategy to tell if the token expires soon
-    return t.expirationTime() < System.currentTimeMillis() - 20_000;
-  }
-
-  private Duration delayBeforeTokenRenewal(Token token) {
-    long expiresIn = token.expirationTime() - System.currentTimeMillis();
-    // TODO use strategy to decide when to renew token
-    long delay = (long) (expiresIn * 0.8);
-    return Duration.ofMillis(delay);
+  private boolean expiresSoon(Token ignores) {
+    return false;
   }
 
   private Token getToken() {
@@ -138,10 +133,7 @@ private void scheduleRenewal(Token t) {
       if (this.renewalTask != null) {
         this.renewalTask.cancel(false);
       }
-      Duration delay = delayBeforeTokenRenewal(t);
-      if (delay.isZero() || delay.isNegative()) {
-        delay = Duration.ofSeconds(1);
-      }
+      Duration delay = this.refreshDelayStrategy.apply(t.expirationTime());
       if (!this.registrations.isEmpty()) {
         LOGGER.debug("Scheduling token retrieval in {}", delay);
         this.renewalTask =
@@ -168,8 +160,8 @@ private void scheduleRenewal(Token t) {
     }
   }
 
-  private static String format(long timestampMs) {
-    return DateTimeFormatter.ISO_INSTANT.format(Instant.ofEpochMilli(timestampMs));
+  private static String format(Instant instant) {
+    return DateTimeFormatter.ISO_INSTANT.format(instant);
   }
 
   private final class RegistrationImpl implements Registration {
@@ -259,4 +251,32 @@ public int hashCode() {
       return Objects.hashCode(id);
     }
   }
+
+  static Function<Instant, Duration> ratioRefreshDelayStrategy(float ratio) {
+    return new RatioRefreshDelayStrategy(ratio);
+  }
+
+  private static class RatioRefreshDelayStrategy implements Function<Instant, Duration> {
+
+    private final float ratio;
+
+    private RatioRefreshDelayStrategy(float ratio) {
+      if (ratio < 0 || ratio > 1) {
+        throw new IllegalArgumentException("Ratio should be > 0 and <= 1: " + ratio);
+      }
+      this.ratio = ratio;
+    }
+
+    @Override
+    public Duration apply(Instant expirationTime) {
+      Duration expiresIn = Duration.between(Instant.now(), expirationTime);
+      Duration delay;
+      if (expiresIn.isZero() || expiresIn.isNegative()) {
+        delay = Duration.ofSeconds(1);
+      } else {
+        delay = Duration.ofMillis((long) (expiresIn.toMillis() * ratio));
+      }
+      return delay;
+    }
+  }
 }

src/main/java/com/rabbitmq/client/amqp/oauth/GsonTokenParser.java

@@ -20,6 +20,7 @@
 import com.google.gson.Gson;
 import com.google.gson.reflect.TypeToken;
 import java.time.Duration;
+import java.time.Instant;
 import java.util.Map;
 
 public class GsonTokenParser implements TokenParser {
@@ -33,16 +34,17 @@ public Token parse(String tokenAsString) {
     String accessToken = (String) tokenAsMap.get("access_token");
     // in seconds, see https://www.rfc-editor.org/rfc/rfc6749#section-5.1
     Duration expiresIn = Duration.ofSeconds(((Number) tokenAsMap.get("expires_in")).longValue());
-    long expirationTime = System.currentTimeMillis() + expiresIn.toMillis();
+    Instant expirationTime =
+        Instant.ofEpochMilli(System.currentTimeMillis() + expiresIn.toMillis());
     return new DefaultTokenInfo(accessToken, expirationTime);
   }
 
   private static final class DefaultTokenInfo implements Token {
 
     private final String value;
-    private final long expirationTime;
+    private final Instant expirationTime;
 
-    private DefaultTokenInfo(String value, long expirationTime) {
+    private DefaultTokenInfo(String value, Instant expirationTime) {
       this.value = value;
       this.expirationTime = expirationTime;
     }
@@ -53,7 +55,7 @@ public String value() {
     }
 
     @Override
-    public long expirationTime() {
+    public Instant expirationTime() {
       return this.expirationTime;
     }
   }

src/main/java/com/rabbitmq/client/amqp/oauth/Token.java

@@ -17,9 +17,11 @@
 // [email protected].
 package com.rabbitmq.client.amqp.oauth;
 
+import java.time.Instant;
+
 public interface Token {
 
   String value();
 
-  long expirationTime();
+  Instant expirationTime();
 }

src/test/java/com/rabbitmq/client/amqp/impl/JwtTestUtils.java

@@ -20,6 +20,7 @@
 import com.google.gson.Gson;
 import com.google.gson.reflect.TypeToken;
 import com.rabbitmq.client.amqp.oauth.Token;
+import java.time.Instant;
 import java.util.Base64;
 import java.util.List;
 import java.util.Map;
@@ -87,8 +88,8 @@ public String value() {
       }
 
       @Override
-      public long expirationTime() {
-        return expirationTime;
+      public Instant expirationTime() {
+        return Instant.ofEpochMilli(expirationTime);
       }
     };
   }

src/test/java/com/rabbitmq/client/amqp/impl/TokenCredentialsTest.java

@@ -20,6 +20,8 @@
 import static com.rabbitmq.client.amqp.impl.Assertions.assertThat;
 import static com.rabbitmq.client.amqp.impl.TestUtils.sync;
 import static com.rabbitmq.client.amqp.impl.TestUtils.waitAtMost;
+import static java.time.Duration.ofMillis;
+import static java.time.Duration.ofSeconds;
 import static java.util.stream.Collectors.toList;
 import static java.util.stream.IntStream.range;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -29,10 +31,12 @@
 import com.rabbitmq.client.amqp.oauth.Token;
 import com.rabbitmq.client.amqp.oauth.TokenRequester;
 import java.time.Duration;
+import java.time.Instant;
 import java.util.List;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -59,13 +63,13 @@ void tearDown() throws Exception {
 
   @Test
   void refreshShouldStopOnceUnregistered() throws InterruptedException {
-    Duration tokenExpiry = Duration.ofMillis(50);
+    Duration tokenExpiry = ofMillis(50);
     AtomicInteger requestCount = new AtomicInteger(0);
     when(this.requester.request())
         .thenAnswer(
             ignored -> {
               requestCount.incrementAndGet();
-              return token("ok", System.currentTimeMillis() + tokenExpiry.toMillis());
+              return token("ok", Instant.now().plus(tokenExpiry));
             });
     TokenCredentials credentials =
         new TokenCredentials(this.requester, this.scheduledExecutorService);
@@ -93,11 +97,11 @@ void refreshShouldStopOnceUnregistered() throws InterruptedException {
 
   @Test
   void severalRegistrationsShouldBeRefreshed() throws InterruptedException {
-    Duration tokenExpiry = Duration.ofMillis(50);
+    Duration tokenExpiry = ofMillis(50);
     Duration waitTime = tokenExpiry.dividedBy(4);
     Duration timeout = tokenExpiry.multipliedBy(20);
     when(this.requester.request())
-        .thenAnswer(ignored -> token("ok", System.currentTimeMillis() + tokenExpiry.toMillis()));
+        .thenAnswer(ignored -> token("ok", Instant.now().plus(tokenExpiry)));
     TokenCredentials credentials =
         new TokenCredentials(this.requester, this.scheduledExecutorService);
     int expectedRefreshCountPerConnection = 3;
@@ -142,15 +146,23 @@ void severalRegistrationsShouldBeRefreshed() throws InterruptedException {
     assertThat(totalRefreshCount).hasValue(refreshCountSnapshot + splitCount * 2);
   }
 
-  private static Token token(String value, long expirationTime) {
+  @Test
+  void refreshDelayStrategy() {
+    Duration diff = ofMillis(100);
+    Function<Instant, Duration> strategy = TokenCredentials.ratioRefreshDelayStrategy(0.8f);
+    assertThat(strategy.apply(Instant.now().plusSeconds(10))).isCloseTo(ofSeconds(8), diff);
+    assertThat(strategy.apply(Instant.now().minusSeconds(10))).isEqualTo(ofSeconds(1));
+  }
+
+  private static Token token(String value, Instant expirationTime) {
     return new Token() {
       @Override
       public String value() {
         return value;
       }
 
       @Override
-      public long expirationTime() {
+      public Instant expirationTime() {
         return expirationTime;
       }
     };

src/test/java/com/rabbitmq/client/amqp/oauth/GsonTokenParserTest.java

@@ -20,10 +20,12 @@
 import static com.rabbitmq.client.amqp.oauth.OAuthTestUtils.sampleJsonToken;
 import static java.time.Duration.ofSeconds;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.within;
 
 import java.time.Duration;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
 import java.util.UUID;
-import org.assertj.core.data.Offset;
 import org.junit.jupiter.api.Test;
 
 public class GsonTokenParserTest {
@@ -38,6 +40,6 @@ void parse() {
     Token token = parser.parse(jsonToken);
     assertThat(token.value()).isEqualTo(accessToken);
     assertThat(token.expirationTime())
-        .isCloseTo(System.currentTimeMillis() + expireIn.toMillis(), Offset.offset(1000L));
+        .isCloseTo(Instant.now().plus(expireIn), within(1, ChronoUnit.SECONDS));
   }
 }

src/test/java/com/rabbitmq/client/amqp/oauth/HttpTokenRequesterTest.java

@@ -29,6 +29,7 @@
 import java.io.IOException;
 import java.io.OutputStream;
 import java.time.Duration;
+import java.time.Instant;
 import java.util.Arrays;
 import java.util.Map;
 import java.util.UUID;
@@ -144,8 +145,8 @@ public String value() {
     }
 
     @Override
-    public long expirationTime() {
-      return 0;
+    public Instant expirationTime() {
+      return Instant.EPOCH;
     }
   }
 }