feat: non-blocking rate limiting on refresh. (#1574)

This further improve the efficient use of futures in the refresh by using the ScheduledExecutor's delay function to
rate limit instead of blocking the ScheduleExecutor's execution thread.

Author: hessjcg

core/pom.xml

@@ -143,6 +143,13 @@
       <version>1.1.5</version>
       <scope>test</scope>
     </dependency>
+    <!-- Using this for the deterministic scheduler class only -->
+    <dependency>
+      <groupId>org.jmock</groupId>
+      <artifactId>jmock</artifactId>
+      <version>2.12.0</version>
+      <scope>test</scope>
+    </dependency>
 
     <dependency>
       <groupId>com.github.jnr</groupId>

core/src/main/java/com/google/cloud/sql/core/AsyncRateLimiter.java

@@ -0,0 +1,81 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * 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.sql.core;
+
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.function.LongSupplier;
+
+/**
+ * A simple, constant-time rate limit calculator. Ensures that there is always at least
+ * delayBetweenAttempts milliseconds between attempts.
+ */
+class AsyncRateLimiter {
+  private long nextOperationTimestamp;
+  private final long delayBetweenAttempts;
+  private final LongSupplier currentTimestampMs;
+
+  /**
+   * Creates a new AsyncRateLimiter uses the System.currentTimeMillis() as the current time.
+   *
+   * @param delayBetweenAttempts the required delay in milliseconds between attempts.
+   */
+  AsyncRateLimiter(long delayBetweenAttempts) {
+    this(delayBetweenAttempts, System::currentTimeMillis);
+  }
+
+  /**
+   * Creates a new AsyncRateLimiter which uses a custom function for the current time.
+   *
+   * @param delayBetweenAttempts the required delay in milliseconds between attempts.
+   * @param currentTimestampMs A function that supplies the current time in milliseconds
+   */
+  AsyncRateLimiter(long delayBetweenAttempts, LongSupplier currentTimestampMs) {
+    this.delayBetweenAttempts = delayBetweenAttempts;
+    this.currentTimestampMs = currentTimestampMs;
+  }
+
+  /**
+   * Returns the number of milliseconds to delay before proceeding with the rate limited operation.
+   * If this returns > 0, the operation must call "acquire" again until it returns 0.
+   */
+  private synchronized long nextDelayMs(long nowTimestampMs) {
+    // allow exactly 1 operation to pass the timestamp.
+    if (nextOperationTimestamp <= nowTimestampMs) {
+      nextOperationTimestamp = nowTimestampMs + delayBetweenAttempts;
+      return 0;
+    }
+
+    return nextOperationTimestamp - nowTimestampMs;
+  }
+
+  /**
+   * Returns a future that will be done when the rate limit has been acquired.
+   *
+   * @param executor the executor to use to schedule future checks for available rate limits.
+   */
+  public ListenableFuture<?> acquireAsync(ScheduledExecutorService executor) {
+    long limit = this.nextDelayMs(currentTimestampMs.getAsLong());
+    if (limit > 0) {
+      return Futures.scheduleAsync(
+          () -> this.acquireAsync(executor), limit, TimeUnit.MILLISECONDS, executor);
+    }
+    return Futures.immediateFuture(null);
+  }
+}

core/src/main/java/com/google/cloud/sql/core/CloudSqlInstance.java

@@ -23,7 +23,6 @@
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningScheduledExecutorService;
-import com.google.common.util.concurrent.RateLimiter;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 import java.io.IOException;
 import java.security.KeyPair;
@@ -56,9 +55,7 @@ class CloudSqlInstance {
   private final CloudSqlInstanceName instanceName;
   private final ListenableFuture<KeyPair> keyPair;
   private final Object instanceDataGuard = new Object();
-
-  @SuppressWarnings("UnstableApiUsage")
-  private final RateLimiter forcedRenewRateLimiter;
+  private final AsyncRateLimiter rateLimiter;
 
   private final RefreshCalculator refreshCalculator = new RefreshCalculator();
 
@@ -89,13 +86,15 @@ class CloudSqlInstance {
       CredentialFactory tokenSourceFactory,
       ListeningScheduledExecutorService executor,
       ListenableFuture<KeyPair> keyPair,
-      @SuppressWarnings("UnstableApiUsage") RateLimiter forcedRenewRateLimiter) {
+      long minRefreshDelayMs) {
     this.instanceName = new CloudSqlInstanceName(connectionName);
     this.instanceDataSupplier = instanceDataSupplier;
     this.authType = authType;
     this.executor = executor;
     this.keyPair = keyPair;
-    this.forcedRenewRateLimiter = forcedRenewRateLimiter;
+
+    // convert requests/second into milliseconds between requests.
+    this.rateLimiter = new AsyncRateLimiter(minRefreshDelayMs);
 
     if (authType == AuthType.IAM) {
       this.accessTokenSupplier = new DefaultAccessTokenSupplier(tokenSourceFactory);
@@ -223,25 +222,19 @@ private ListenableFuture<InstanceData> startRefreshAttempt() {
       refreshRunning = true;
     }
 
-    // To avoid unreasonable SQL Admin API usage, use a rate limit to throttle our usage.
-    ListenableFuture<?> rateLimit =
-        executor.submit(
-            () -> {
-              logger.fine(
-                  String.format(
-                      "[%s] Refresh Operation: Acquiring rate limiter permit.", instanceName));
-              //noinspection UnstableApiUsage
-              forcedRenewRateLimiter.acquire();
-              logger.fine(
-                  String.format(
-                      "[%s] Refresh Operation: Acquired rate limiter permit. Starting refresh...",
-                      instanceName));
-            },
-            executor);
+    logger.fine(
+        String.format("[%s] Refresh Operation: Acquiring rate limiter permit.", instanceName));
+    ListenableFuture<?> delay = rateLimiter.acquireAsync(executor);
+    delay.addListener(
+        () ->
+            logger.fine(
+                String.format(
+                    "[%s] Refresh Operation: Rate limiter permit acquired.", instanceName)),
+        executor);
 
     // Once rate limiter is done, attempt to getInstanceData.
     ListenableFuture<InstanceData> dataFuture =
-        Futures.whenAllComplete(rateLimit)
+        Futures.whenAllComplete(delay)
             .callAsync(
                 () ->
                     instanceDataSupplier.getInstanceData(

core/src/main/java/com/google/cloud/sql/core/CoreSocketFactory.java

@@ -25,7 +25,6 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningScheduledExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
-import com.google.common.util.concurrent.RateLimiter;
 import java.io.File;
 import java.io.IOException;
 import java.net.InetSocketAddress;
@@ -108,6 +107,7 @@ public final class CoreSocketFactory {
   private static final int RSA_KEY_SIZE = 2048;
   private static final List<String> userAgents = new ArrayList<>();
   private static final String version = getVersion();
+  private static final long MIN_REFRESH_DELAY_MS = 30000; // Minimum 30 seconds between refresh.
   private static CoreSocketFactory coreSocketFactory;
   private final ListenableFuture<KeyPair> localKeyPair;
   private final ConcurrentHashMap<String, CloudSqlInstance> instances = new ConcurrentHashMap<>();
@@ -392,6 +392,6 @@ private CloudSqlInstance apiFetcher(ConnectionConfig config) {
         instanceCredentialFactory,
         executor,
         localKeyPair,
-        RateLimiter.create(1.0 / 30.0)); // 1 refresh attempt every 30 seconds
+        MIN_REFRESH_DELAY_MS);
   }
 }

core/src/test/java/com/google/cloud/sql/core/AsyncRateLimiterTest.java

@@ -0,0 +1,137 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * 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.sql.core;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import org.jmock.lib.concurrent.DeterministicScheduler;
+import org.junit.Test;
+
+public class AsyncRateLimiterTest {
+
+  @Test
+  public void firstCallShouldReturnNoDelay() {
+    RateLimiterTestHarness th = new RateLimiterTestHarness(100);
+    ListenableFuture<?> f = th.rateLimiter.acquireAsync(th.ex);
+    assertThat(f.isDone()).isTrue();
+  }
+
+  @Test
+  public void subsequentCallsShouldReturnDelays() {
+    RateLimiterTestHarness th = new RateLimiterTestHarness(100);
+    ListenableFuture<?> f1 = th.rateLimiter.acquireAsync(th.ex);
+    ListenableFuture<?> f2 = th.rateLimiter.acquireAsync(th.ex);
+    // When calls occur at the same timestamp, one will return 0 and
+    // the other will return the full delay.
+    assertThat(f1.isDone()).isTrue();
+    assertThat(f2.isDone()).isFalse();
+  }
+
+  @Test
+  public void delayBeforeExpiration() throws InterruptedException {
+    RateLimiterTestHarness th = new RateLimiterTestHarness(100);
+    ListenableFuture<?> f1 = th.rateLimiter.acquireAsync(th.ex);
+    assertThat(f1.isDone()).isTrue();
+
+    // Before the expiration, isDone should be false
+    th.tickMs(99);
+    ListenableFuture<?> f2 = th.rateLimiter.acquireAsync(th.ex);
+    assertThat(f2.isDone()).isFalse();
+
+    // Exactly at the expiration, isDone should be true
+    th.tickMs(1);
+    assertThat(f2.isDone()).isTrue();
+  }
+
+  @Test
+  public void noDelayExactlyAtExpiration() throws InterruptedException {
+    RateLimiterTestHarness th = new RateLimiterTestHarness(100);
+    ListenableFuture<?> f1 = th.rateLimiter.acquireAsync(th.ex);
+    ListenableFuture<?> f2 = th.rateLimiter.acquireAsync(th.ex);
+    assertThat(f1.isDone()).isTrue();
+    assertThat(f2.isDone()).isFalse();
+
+    th.tickMs(50);
+    assertThat(f2.isDone()).isFalse();
+
+    th.tickMs(100);
+    assertThat(f2.isDone()).isTrue();
+  }
+
+  @Test
+  public void noDelayAfterExpiration() throws InterruptedException {
+    RateLimiterTestHarness th = new RateLimiterTestHarness(100);
+    ListenableFuture<?> f1 = th.rateLimiter.acquireAsync(th.ex);
+    assertThat(f1.isDone()).isTrue();
+
+    th.tickMs(101);
+    ListenableFuture<?> f2 = th.rateLimiter.acquireAsync(th.ex);
+    assertThat(f2.isDone()).isTrue();
+  }
+
+  @Test
+  public void testAsyncWorks() {
+    final long delay = 100;
+
+    RateLimiterTestHarness th = new RateLimiterTestHarness(delay);
+
+    List<ListenableFuture<?>> futures = new ArrayList<>();
+    for (int i = 0; i < 3; i++) {
+      futures.add(th.rateLimiter.acquireAsync(th.ex));
+    }
+
+    // First attempt happens without any delay, because nothing is run yet.
+    assertThat(futures.stream().mapToInt(f -> f.isDone() ? 1 : 0).sum()).isEqualTo(1);
+
+    // Tick forward less than the delay.
+    th.tickMs(50);
+
+    // When all futures are evaluated again immediately, still only 1 request has finished.
+    // because not enough time has elapsed.
+    assertThat(futures.stream().mapToInt(f -> f.isDone() ? 1 : 0).sum()).isEqualTo(1);
+
+    // Tick forward more than the delay. Now 2 attempts should have finished.
+    th.tickMs(100);
+
+    assertThat(futures.stream().mapToInt(f -> f.isDone() ? 1 : 0).sum()).isEqualTo(2);
+
+    // Tick forward more than the delay. Now 3 attempts should have finished.
+    th.tickMs(100);
+    assertThat(futures.stream().mapToInt(f -> f.isDone() ? 1 : 0).sum()).isEqualTo(3);
+  }
+
+  private static class RateLimiterTestHarness {
+
+    final AtomicLong now = new AtomicLong(System.currentTimeMillis());
+    final DeterministicScheduler ex = new DeterministicScheduler();
+    final AsyncRateLimiter rateLimiter;
+
+    RateLimiterTestHarness(long delay) {
+      rateLimiter = new AsyncRateLimiter(delay, now::get);
+    }
+
+    private void tickMs(long ms) {
+      now.addAndGet(ms);
+      ex.tick(ms, TimeUnit.MILLISECONDS);
+    }
+  }
+}

core/src/test/java/com/google/cloud/sql/core/CloudSqlInstanceConcurrencyTest.java

@@ -25,7 +25,6 @@
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningScheduledExecutorService;
-import com.google.common.util.concurrent.RateLimiter;
 import java.io.IOException;
 import java.security.KeyPair;
 import java.util.ArrayList;
@@ -56,6 +55,7 @@ public void initialize(HttpRequest var1) throws IOException {
 
   @Test(timeout = 20000) // 45 seconds timeout in case of deadlock
   public void testForceRefreshDoesNotCauseADeadlockOrBrokenRefreshLoop() throws Exception {
+    final long refreshDelayMs = 50;
     MockAdminApi mockAdminApi = new MockAdminApi();
     ListenableFuture<KeyPair> keyPairFuture =
         Futures.immediateFuture(mockAdminApi.getClientKeyPair());
@@ -74,7 +74,7 @@ public void testForceRefreshDoesNotCauseADeadlockOrBrokenRefreshLoop() throws Ex
               new TestCredentialFactory(),
               executor,
               keyPairFuture,
-              newRateLimiter()));
+              refreshDelayMs));
     }
 
     // Get SSL Data for each instance, forcing the first refresh to complete.
@@ -130,8 +130,4 @@ private Thread startForceRefreshThread(CloudSqlInstance inst) {
     t.start();
     return t;
   }
-
-  private RateLimiter newRateLimiter() {
-    return RateLimiter.create(20.0); // 20/sec = every 50 ms
-  }
 }