test(storage): retry transient socket connections in testbench control plane

Author: Dhriti Chopra

java-storage/google-cloud-storage/src/test/java/com/google/cloud/storage/it/runner/registry/TestBench.java

@@ -50,6 +50,7 @@
 import java.io.InputStreamReader;
 import java.net.SocketException;
 import java.net.URI;
+import java.util.concurrent.Callable;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -150,7 +151,7 @@ public RetryTestResource createRetryTest(RetryTestResource retryTestResource) th
     HttpContent content =
         new ByteArrayContent("application/json", jsonString.getBytes(StandardCharsets.UTF_8));
     HttpRequest req = requestFactory.buildPostRequest(url, content);
-    HttpResponse resp = req.execute();
+    HttpResponse resp = execWithRetries(req::execute);
     RetryTestResource result = gson.fromJson(resp.parseAsString(), RetryTestResource.class);
     resp.disconnect();
     return result;
@@ -159,14 +160,14 @@ public RetryTestResource createRetryTest(RetryTestResource retryTestResource) th
   public void deleteRetryTest(RetryTestResource retryTestResource) throws IOException {
     GenericUrl url = new GenericUrl(baseUri + "/retry_test/" + retryTestResource.id);
     HttpRequest req = requestFactory.buildDeleteRequest(url);
-    HttpResponse resp = req.execute();
+    HttpResponse resp = execWithRetries(req::execute);
     resp.disconnect();
   }
 
   public RetryTestResource getRetryTest(RetryTestResource retryTestResource) throws IOException {
     GenericUrl url = new GenericUrl(baseUri + "/retry_test/" + retryTestResource.id);
     HttpRequest req = requestFactory.buildGetRequest(url);
-    HttpResponse resp = req.execute();
+    HttpResponse resp = execWithRetries(req::execute);
     RetryTestResource result = gson.fromJson(resp.parseAsString(), RetryTestResource.class);
     resp.disconnect();
     return result;
@@ -175,7 +176,7 @@ public RetryTestResource getRetryTest(RetryTestResource retryTestResource) throw
   public List<RetryTestResource> listRetryTests() throws IOException {
     GenericUrl url = new GenericUrl(baseUri + "/retry_tests");
     HttpRequest req = requestFactory.buildGetRequest(url);
-    HttpResponse resp = req.execute();
+    HttpResponse resp = execWithRetries(req::execute);
     JsonObject result = gson.fromJson(resp.parseAsString(), JsonObject.class);
     JsonArray retryTest = (JsonArray) result.get("retry_test");
     ImmutableList.Builder<RetryTestResource> b = ImmutableList.builder();
@@ -186,6 +187,33 @@ public List<RetryTestResource> listRetryTests() throws IOException {
     return b.build();
   }
 
+  private <T> T execWithRetries(Callable<T> callable) throws IOException {
+    try {
+      return runWithRetries(
+          callable,
+          RetrySettings.newBuilder()
+              .setTotalTimeoutDuration(Duration.ofSeconds(15))
+              .setInitialRetryDelayDuration(Duration.ofMillis(200))
+              .setRetryDelayMultiplier(1.5)
+              .setMaxRetryDelayDuration(Duration.ofSeconds(2))
+              .setMaxAttempts(5)
+              .build(),
+          new BasicResultRetryAlgorithm<T>() {
+            @Override
+            public boolean shouldRetry(Throwable previousThrowable, T previousResponse) {
+              return previousThrowable instanceof SocketException
+                  || previousThrowable instanceof IOException;
+            }
+          },
+          NanoClock.getDefaultClock());
+    } catch (RetryHelperException e) {
+      if (e.getCause() instanceof IOException) {
+        throw (IOException) e.getCause();
+      }
+      throw new IOException(e);
+    }
+  }
+
   private boolean startGRPCServer(int gRPCPort) throws IOException {
     GenericUrl url = new GenericUrl(baseUri + "/start_grpc?port=9090");
     HttpRequest req = requestFactory.buildGetRequest(url);