fix: re-use existing connection info during force refresh (#180)

GoogleCloudPlatform · Sep 18, 2023 · 75d667a · 75d667a
1 parent d72f517
commit 75d667a
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 7 deletions.
diff --git a/...ydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/DefaultConnectionInfoCache.java b/...ydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/DefaultConnectionInfoCache.java
@@ -187,8 +187,7 @@ public void forceRefresh() {
               "[%s] Force Refresh: the next refresh operation was cancelled."
                   + " Scheduling new refresh operation immediately.",
               instanceName));
-      current = executor.submit(this::performRefresh);
-      next = current;
+      next = executor.submit(this::performRefresh);
     }
   }
 }
diff --git a/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ConnectionInfoCacheTest.java b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ConnectionInfoCacheTest.java
@@ -41,6 +41,7 @@
 
 public class ConnectionInfoCacheTest {
 
+  private static final int DEFAULT_WAIT = 100;
   private static final String TEST_INSTANCE_IP = "10.0.0.1";
   private static final String TEST_INSTANCE_ID = "some-instance-id";
   private static final Instant ONE_HOUR_FROM_NOW = Instant.now().plus(1, ChronoUnit.HOURS);
@@ -288,7 +289,7 @@ public void testGetConnectionInfo_isRateLimited() {
   }
 
   @Test
-  public void testForceRefresh_schedulesNextRefreshImmediately() {
+  public void testForceRefresh_schedulesNextRefreshImmediately() throws InterruptedException {
     ScheduledExecutorService executor = Executors.newScheduledThreadPool(2);
 
     InMemoryConnectionInfoRepo connectionInfoRepo = new InMemoryConnectionInfoRepo();
@@ -329,8 +330,20 @@ public void testForceRefresh_schedulesNextRefreshImmediately() {
 
     connectionInfoCache.forceRefresh();
 
-    // After the force refresh, new refresh data is available.
+    // Refresh data hasn't changed because we re-use the existing connection info.
     connectionInfo = connectionInfoCache.getConnectionInfo();
+    assertThat(connectionInfoRepo.getIndex()).isEqualTo(1);
+
+    for (int i = 0; i < 10; i++) {
+      connectionInfo = connectionInfoCache.getConnectionInfo();
+      if (connectionInfoRepo.getIndex() > 1) {
+        break;
+      }
+      Thread.sleep(DEFAULT_WAIT);
+    }
+
+    // After a few seconds, new refresh data should be available.
+    assertThat(connectionInfoRepo.getIndex()).isEqualTo(2);
     assertThat(
             connectionInfo
                 .getClientCertificate()
@@ -341,7 +354,8 @@ public void testForceRefresh_schedulesNextRefreshImmediately() {
   }
 
   @Test
-  public void testForceRefresh_refreshCalledOnlyOnceDuringMultipleCalls() {
+  public void testForceRefresh_refreshCalledOnlyOnceDuringMultipleCalls()
+      throws InterruptedException {
     ScheduledExecutorService executor = Executors.newScheduledThreadPool(2);
 
     InMemoryConnectionInfoRepo connectionInfoRepo = new InMemoryConnectionInfoRepo();
@@ -389,9 +403,18 @@ public void testForceRefresh_refreshCalledOnlyOnceDuringMultipleCalls() {
     connectionInfoCache.forceRefresh();
     // This second call should be ignored as there is a refresh operation in progress.
     connectionInfoCache.forceRefresh();
+    assertThat(connectionInfoRepo.getIndex()).isEqualTo(1);
+
+    for (int i = 0; i < 10; i++) {
+      connectionInfo = connectionInfoCache.getConnectionInfo();
+      if (connectionInfoRepo.getIndex() > 1) {
+        break;
+      }
+      Thread.sleep(DEFAULT_WAIT);
+    }
 
-    // After the force refresh, new refresh data is available.
-    connectionInfo = connectionInfoCache.getConnectionInfo();
+    // After a few seconds, new refresh data should be available.
+    assertThat(connectionInfoRepo.getIndex()).isEqualTo(2);
     assertThat(
             connectionInfo
                 .getClientCertificate()

diff --git a/...ydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/InMemoryConnectionInfoRepo.java b/...ydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/InMemoryConnectionInfoRepo.java
@@ -55,4 +55,8 @@ public ConnectionInfo getConnectionInfo(InstanceName instanceName, KeyPair publi
   public final void addResponses(Callable<ConnectionInfo>... callables) {
     registeredCallables.addAll(Arrays.asList(callables));
   }
+
+  public final int getIndex() {
+    return this.index.get();
+  }
 }