From 75d667a6037d8dc69b7a06f15df93895deac0bf1 Mon Sep 17 00:00:00 2001
From: Tatiane Tosta <91583351+ttosta-google@users.noreply.github.com>
Date: Mon, 18 Sep 2023 17:57:01 +0000
Subject: [PATCH] fix: re-use existing connection info during force refresh
 (#180)

---
 .../alloydb/DefaultConnectionInfoCache.java   |  3 +-
 .../alloydb/ConnectionInfoCacheTest.java      | 33 ++++++++++++++++---
 .../alloydb/InMemoryConnectionInfoRepo.java   |  4 +++
 3 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/DefaultConnectionInfoCache.java b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/DefaultConnectionInfoCache.java
index f3c74144..f152b8c2 100644
--- a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/DefaultConnectionInfoCache.java
+++ b/alloydb-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
index a7346018..0cce9c68 100644
--- 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/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/InMemoryConnectionInfoRepo.java b/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/InMemoryConnectionInfoRepo.java
index 399d5d66..95194c6e 100644
--- a/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/InMemoryConnectionInfoRepo.java
+++ b/alloydb-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();
+  }
 }