chore: Remove ListenableFuture from SqlAdminApiFetcher (#1307)

hessjcg · web-flow · commit 01d5036c3c6c · 2023-06-08T10:52:10.000-06:00
Changes the return type of SqlAdminApiFetcher.getInstanceData() from ListenableFuture&lt;InstanceData&gt; to
InstanceData in an effort to simplify Future related code across this project.
diff --git a/core/pom.xml b/core/pom.xml
@@ -211,6 +211,12 @@
       <version>1.13.0</version>
     </dependency>
 
+    <dependency>
+      <groupId>com.google.auth</groupId>
+      <artifactId>google-auth-library-credentials</artifactId>
+      <version>1.13.0</version>
+    </dependency>
+
     <dependency>
       <groupId>com.google.oauth-client</groupId>
       <artifactId>google-oauth-client</artifactId>
diff --git a/core/src/main/java/com/google/cloud/sql/core/CloudSqlInstance.java b/core/src/main/java/com/google/cloud/sql/core/CloudSqlInstance.java
@@ -25,7 +25,6 @@
 import com.google.cloud.sql.AuthType;
 import com.google.cloud.sql.CredentialFactory;
 import com.google.common.base.Throwables;
-import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningScheduledExecutorService;
@@ -42,7 +41,6 @@
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
-import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.net.ssl.SSLSocket;
 
@@ -89,8 +87,7 @@ class CloudSqlInstance {
       AuthType authType,
       CredentialFactory tokenSourceFactory,
       ListeningScheduledExecutorService executor,
-      ListenableFuture<KeyPair> keyPair)
-      throws IOException, InterruptedException {
+      ListenableFuture<KeyPair> keyPair) {
     this.instanceName = new CloudSqlInstanceName(connectionName);
     this.apiFetcher = apiFetcher;
     this.authType = authType;
@@ -228,61 +225,38 @@ private InstanceData performRefresh() throws InterruptedException, ExecutionExce
     // To avoid unreasonable SQL Admin API usage, use a rate limit to throttle our usage.
     forcedRenewRateLimiter.acquirePermit();
 
-    ListenableFuture<InstanceData> refreshFuture;
-
+    final GoogleCredentials credentials;
     if (iamAuthnCredentials.isPresent()) {
-      GoogleCredentials downscopedCredentials = getDownscopedCredentials(iamAuthnCredentials.get());
-      refreshFuture =
-          apiFetcher.getInstanceData(
-              instanceName, downscopedCredentials, AuthType.IAM, executor, keyPair);
+      credentials = getDownscopedCredentials(iamAuthnCredentials.get());
     } else {
-      refreshFuture =
-          apiFetcher.getInstanceData(instanceName, null, AuthType.PASSWORD, executor, keyPair);
+      credentials = null;
     }
 
-    Futures.addCallback(
-        refreshFuture,
-        new FutureCallback<InstanceData>() {
-          @Override
-          public void onSuccess(InstanceData instanceData) {
-            synchronized (instanceDataGuard) {
-              // update currentInstanceData with the most recent results
-              currentInstanceData = refreshFuture;
-              // schedule a replacement before the SSLContext expires;
-              nextInstanceData =
-                  executor.schedule(
-                      () -> performRefresh(),
-                      refreshCalculator.calculateSecondsUntilNextRefresh(
-                          Instant.now(), getInstanceData().getExpiration().toInstant()),
-                      TimeUnit.SECONDS);
-            }
-          }
+    try {
+      InstanceData data =
+          apiFetcher.getInstanceData(
+              this.instanceName, credentials, this.authType, executor, keyPair);
+
+      synchronized (instanceDataGuard) {
+        // update currentInstanceData with the most recent results
+        currentInstanceData = Futures.immediateFuture(data);
 
-          @Override
-          public void onFailure(Throwable t) {
-            logger.log(
-                Level.WARNING,
-                "An error occurred while performing refresh. Retrying immediately.",
-                t);
-            synchronized (instanceDataGuard) {
-              InstanceData instanceData = null;
-              try {
-                instanceData = getInstanceData();
-              } catch (Exception e) {
-                // this means the result was invalid
-              }
-              if (instanceData == null
-                  || instanceData.getExpiration().toInstant().isBefore(Instant.now())) {
-                // replace current if it is expired or invalid
-                currentInstanceData = refreshFuture;
-              }
-              nextInstanceData = executor.submit(() -> performRefresh());
-            }
-          }
-        },
-        executor);
+        // schedule a replacement before the SSLContext expires;
+        nextInstanceData =
+            executor.schedule(
+                this::performRefresh,
+                refreshCalculator.calculateSecondsUntilNextRefresh(
+                    Instant.now(), data.getExpiration().toInstant()),
+                TimeUnit.SECONDS);
+      }
 
-    return refreshFuture.get();
+      return data;
+    } catch (ExecutionException | InterruptedException e) {
+      synchronized (instanceDataGuard) {
+        nextInstanceData = executor.submit(this::performRefresh);
+      }
+      throw e;
+    }
   }
 
   private Optional<Date> getTokenExpirationTime(Credential credentials) {
diff --git a/core/src/main/java/com/google/cloud/sql/core/CoreSocketFactory.java b/core/src/main/java/com/google/cloud/sql/core/CoreSocketFactory.java
@@ -355,13 +355,8 @@ Socket createSslSocket(String instanceName, List<String> ipTypes)
   CloudSqlInstance getCloudSqlInstance(String instanceName, AuthType authType) {
     return instances.computeIfAbsent(
         instanceName,
-        k -> {
-          try {
-            return new CloudSqlInstance(
-                k, adminApiService, authType, credentialFactory, executor, localKeyPair);
-          } catch (IOException | InterruptedException e) {
-            throw new RuntimeException(e);
-          }
-        });
+        k ->
+            new CloudSqlInstance(
+                k, adminApiService, authType, credentialFactory, executor, localKeyPair));
   }
 }
diff --git a/core/src/main/java/com/google/cloud/sql/core/SqlAdminApiFetcher.java b/core/src/main/java/com/google/cloud/sql/core/SqlAdminApiFetcher.java
@@ -53,6 +53,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
 import java.util.logging.Logger;
 import javax.net.ssl.KeyManagerFactory;
 import javax.net.ssl.SSLContext;
@@ -93,12 +94,14 @@ private String generatePublicKeyCert(KeyPair keyPair) {
         + "-----END RSA PUBLIC KEY-----\n";
   }
 
-  ListenableFuture<InstanceData> getInstanceData(
+  InstanceData getInstanceData(
       CloudSqlInstanceName instanceName,
       OAuth2Credentials credentials,
       AuthType authType,
       ListeningScheduledExecutorService executor,
-      ListenableFuture<KeyPair> keyPair) {
+      ListenableFuture<KeyPair> keyPair)
+      throws ExecutionException, InterruptedException {
+
     // Fetch the metadata
     ListenableFuture<Metadata> metadataFuture =
         executor.submit(() -> fetchMetadata(instanceName, authType));
@@ -152,7 +155,7 @@ ListenableFuture<InstanceData> getInstanceData(
                 },
                 executor);
 
-    return done;
+    return done.get();
   }
 
   private Optional<Date> getTokenExpirationTime(OAuth2Credentials credentials) {
diff --git a/core/src/test/java/com/google/cloud/sql/core/BadConnectionFactory.java b/core/src/test/java/com/google/cloud/sql/core/BadConnectionFactory.java
@@ -38,6 +38,11 @@ public void addHeader(String name, String value) throws IOException {
 
     @Override
     public LowLevelHttpResponse execute() throws IOException {
+      try {
+        Thread.sleep(100);
+      } catch (InterruptedException e) {
+        // Ignore the interruption
+      }
       throw new SocketTimeoutException("Fake connect timeout");
     }
   }
diff --git a/core/src/test/java/com/google/cloud/sql/core/SqlAdminApiFetcherTest.java b/core/src/test/java/com/google/cloud/sql/core/SqlAdminApiFetcherTest.java
@@ -23,7 +23,6 @@
 import com.google.auth.oauth2.OAuth2CredentialsWithRefresh;
 import com.google.cloud.sql.AuthType;
 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.MoreExecutors;
 import java.security.GeneralSecurityException;
@@ -55,15 +54,13 @@ public void testFetchInstanceData_returnsIpAddresses()
         new StubApiFetcherFactory(mockAdminApi.getHttpTransport())
             .create(new StubCredentialFactory().create());
 
-    ListenableFuture<InstanceData> instanceDataFuture =
+    InstanceData instanceData =
         fetcher.getInstanceData(
             new CloudSqlInstanceName(INSTANCE_CONNECTION_NAME),
             null,
             AuthType.PASSWORD,
             newTestExecutor(),
             Futures.immediateFuture(mockAdminApi.getClientKeyPair()));
-
-    InstanceData instanceData = instanceDataFuture.get();
     assertThat(instanceData.getSslContext()).isInstanceOf(SSLContext.class);
 
     Map<String, String> ipAddrs = instanceData.getIpAddrs();
@@ -89,20 +86,22 @@ public void testFetchInstanceData_throwsException_whenIamAuthnIsNotSupported()
         new StubApiFetcherFactory(mockAdminApi.getHttpTransport())
             .create(new StubCredentialFactory().create());
 
-    ListenableFuture<InstanceData> instanceData =
-        fetcher.getInstanceData(
-            new CloudSqlInstanceName(INSTANCE_CONNECTION_NAME),
-            OAuth2CredentialsWithRefresh.newBuilder()
-                .setRefreshHandler(
-                    mockAdminApi.getRefreshHandler(
-                        "refresh-token", Date.from(Instant.now().plus(1, ChronoUnit.HOURS))))
-                .setAccessToken(new AccessToken("my-token", Date.from(Instant.now())))
-                .build(),
-            AuthType.IAM,
-            newTestExecutor(),
-            Futures.immediateFuture(mockAdminApi.getClientKeyPair()));
-
-    ExecutionException ex = assertThrows(ExecutionException.class, instanceData::get);
+    ExecutionException ex =
+        assertThrows(
+            ExecutionException.class,
+            () -> {
+              fetcher.getInstanceData(
+                  new CloudSqlInstanceName(INSTANCE_CONNECTION_NAME),
+                  OAuth2CredentialsWithRefresh.newBuilder()
+                      .setRefreshHandler(
+                          mockAdminApi.getRefreshHandler(
+                              "refresh-token", Date.from(Instant.now().plus(1, ChronoUnit.HOURS))))
+                      .setAccessToken(new AccessToken("my-token", Date.from(Instant.now())))
+                      .build(),
+                  AuthType.IAM,
+                  newTestExecutor(),
+                  Futures.immediateFuture(mockAdminApi.getClientKeyPair()));
+            });
     assertThat(ex)
         .hasMessageThat()
         .contains("[p:r:i] IAM Authentication is not supported for SQL Server instances");
@@ -116,20 +115,22 @@ public void testFetchInstanceData_throwsException_whenTokenIsEmpty()
         new StubApiFetcherFactory(mockAdminApi.getHttpTransport())
             .create(new StubCredentialFactory().create());
 
-    ListenableFuture<InstanceData> instanceData =
-        fetcher.getInstanceData(
-            new CloudSqlInstanceName(INSTANCE_CONNECTION_NAME),
-            OAuth2CredentialsWithRefresh.newBuilder()
-                .setRefreshHandler(
-                    mockAdminApi.getRefreshHandler(
-                        "", Date.from(Instant.now().plus(1, ChronoUnit.HOURS)) /* empty */))
-                .setAccessToken(new AccessToken("" /* ignored */, Date.from(Instant.now())))
-                .build(),
-            AuthType.IAM,
-            newTestExecutor(),
-            Futures.immediateFuture(mockAdminApi.getClientKeyPair()));
-
-    ExecutionException ex = assertThrows(ExecutionException.class, instanceData::get);
+    ExecutionException ex =
+        assertThrows(
+            ExecutionException.class,
+            () -> {
+              fetcher.getInstanceData(
+                  new CloudSqlInstanceName(INSTANCE_CONNECTION_NAME),
+                  OAuth2CredentialsWithRefresh.newBuilder()
+                      .setRefreshHandler(
+                          mockAdminApi.getRefreshHandler(
+                              "", Date.from(Instant.now().plus(1, ChronoUnit.HOURS)) /* empty */))
+                      .setAccessToken(new AccessToken("" /* ignored */, Date.from(Instant.now())))
+                      .build(),
+                  AuthType.IAM,
+                  newTestExecutor(),
+                  Futures.immediateFuture(mockAdminApi.getClientKeyPair()));
+            });
 
     assertThat(ex).hasMessageThat().contains("Access Token has length of zero");
   }
@@ -142,21 +143,23 @@ public void testFetchInstanceData_throwsException_whenTokenIsExpired()
         new StubApiFetcherFactory(mockAdminApi.getHttpTransport())
             .create(new StubCredentialFactory().create());
 
-    ListenableFuture<InstanceData> instanceData =
-        fetcher.getInstanceData(
-            new CloudSqlInstanceName(INSTANCE_CONNECTION_NAME),
-            OAuth2CredentialsWithRefresh.newBuilder()
-                .setRefreshHandler(
-                    mockAdminApi.getRefreshHandler(
-                        "refresh-token",
-                        Date.from(Instant.now().minus(1, ChronoUnit.HOURS)) /* 1 hour ago */))
-                .setAccessToken(new AccessToken("original-token", Date.from(Instant.now())))
-                .build(),
-            AuthType.IAM,
-            newTestExecutor(),
-            Futures.immediateFuture(mockAdminApi.getClientKeyPair()));
-
-    ExecutionException ex = assertThrows(ExecutionException.class, instanceData::get);
+    ExecutionException ex =
+        assertThrows(
+            ExecutionException.class,
+            () -> {
+              fetcher.getInstanceData(
+                  new CloudSqlInstanceName(INSTANCE_CONNECTION_NAME),
+                  OAuth2CredentialsWithRefresh.newBuilder()
+                      .setRefreshHandler(
+                          mockAdminApi.getRefreshHandler(
+                              "refresh-token",
+                              Date.from(Instant.now().minus(1, ChronoUnit.HOURS)) /* 1 hour ago */))
+                      .setAccessToken(new AccessToken("original-token", Date.from(Instant.now())))
+                      .build(),
+                  AuthType.IAM,
+                  newTestExecutor(),
+                  Futures.immediateFuture(mockAdminApi.getClientKeyPair()));
+            });
 
     assertThat(ex).hasMessageThat().contains("Access Token expiration time is in the past");
   }
@@ -169,21 +172,24 @@ public void testFetchInstanceData_throwsException_whenRequestsTimeout()
         new StubApiFetcherFactory(new BadConnectionFactory())
             .create(new StubCredentialFactory().create());
 
-    ListenableFuture<InstanceData> instanceData =
-        fetcher.getInstanceData(
-            new CloudSqlInstanceName(INSTANCE_CONNECTION_NAME),
-            OAuth2CredentialsWithRefresh.newBuilder()
-                .setRefreshHandler(
-                    mockAdminApi.getRefreshHandler(
-                        "refresh-token",
-                        Date.from(Instant.now().plus(1, ChronoUnit.HOURS)) /* 1 hour from now */))
-                .setAccessToken(new AccessToken("original-token", Date.from(Instant.now())))
-                .build(),
-            AuthType.IAM,
-            newTestExecutor(),
-            Futures.immediateFuture(mockAdminApi.getClientKeyPair()));
-
-    ExecutionException ex = assertThrows(ExecutionException.class, instanceData::get);
+    ExecutionException ex =
+        assertThrows(
+            ExecutionException.class,
+            () -> {
+              fetcher.getInstanceData(
+                  new CloudSqlInstanceName(INSTANCE_CONNECTION_NAME),
+                  OAuth2CredentialsWithRefresh.newBuilder()
+                      .setRefreshHandler(
+                          mockAdminApi.getRefreshHandler(
+                              "refresh-token",
+                              Date.from(
+                                  Instant.now().plus(1, ChronoUnit.HOURS)) /* 1 hour from now */))
+                      .setAccessToken(new AccessToken("original-token", Date.from(Instant.now())))
+                      .build(),
+                  AuthType.IAM,
+                  newTestExecutor(),
+                  Futures.immediateFuture(mockAdminApi.getClientKeyPair()));
+            });
 
     assertThat(ex.getCause().getCause()).hasMessageThat().contains("Fake connect timeout");
   }
diff --git a/core/src/test/java/com/google/cloud/sql/core/StubCredentialFactory.java b/core/src/test/java/com/google/cloud/sql/core/StubCredentialFactory.java
diff --git a/pom.xml b/pom.xml

Original file line number	Diff line number	Diff line change
`@@ -38,6 +38,11 @@ public void addHeader(String name, String value) throws IOException {`
`38`	`38`
`39`	`39`	`@Override`
`40`	`40`	`public LowLevelHttpResponse execute() throws IOException {`
	`41`	`+ try {`
	`42`	`+ Thread.sleep(100);`
	`43`	`+ } catch (InterruptedException e) {`
	`44`	`+ // Ignore the interruption`
	`45`	`+ }`
`41`	`46`	`throw new SocketTimeoutException("Fake connect timeout");`
`42`	`47`	`}`
`43`	`48`	`}`