Try and simplify shouldRetry override

Author: nicktindall

modules/repository-gcs/src/main/java/org/elasticsearch/repositories/gcs/DelegatingStorageRetryStrategy.java

@@ -14,6 +14,7 @@
 import com.google.api.client.http.HttpTransport;
 import com.google.api.client.http.javanet.NetHttpTransport;
 import com.google.api.client.util.SecurityUtils;
+import com.google.api.gax.retrying.ResultRetryAlgorithm;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auth.oauth2.ServiceAccountCredentials;
 import com.google.cloud.ServiceOptions;
@@ -43,7 +44,6 @@
 import java.net.UnknownHostException;
 import java.security.KeyStore;
 import java.util.Map;
-import java.util.concurrent.CancellationException;
 import java.util.stream.Collectors;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -258,17 +258,14 @@ StorageOptions createStorageOptions(
     }
 
     protected StorageRetryStrategy getRetryStrategy() {
-        return new DelegatingStorageRetryStrategy<>(
+        return ShouldRetryDecorator.decorate(
             StorageRetryStrategy.getLegacyStorageRetryStrategy(),
-            d -> new DelegatingStorageRetryStrategy.DelegatingResultRetryAlgorithm<>(d) {
-                @Override
-                public boolean shouldRetry(Throwable prevThrowable, Object prevResponse) throws CancellationException {
-                    // Retry in the event of an unknown host exception
-                    if (ExceptionsHelper.unwrap(prevThrowable, UnknownHostException.class) != null) {
-                        return true;
-                    }
-                    return delegate.shouldRetry(prevThrowable, prevResponse);
+            (Throwable prevThrowable, Object prevResponse, ResultRetryAlgorithm<Object> delegate) -> {
+                // Retry in the event of an unknown host exception
+                if (ExceptionsHelper.unwrap(prevThrowable, UnknownHostException.class) != null) {
+                    return true;
                 }
+                return delegate.shouldRetry(prevThrowable, prevResponse);
             }
         );
     }

modules/repository-gcs/src/main/java/org/elasticsearch/repositories/gcs/GoogleCloudStorageService.java

@@ -0,0 +1,81 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.repositories.gcs;
+
+import com.google.api.gax.retrying.ResultRetryAlgorithm;
+import com.google.api.gax.retrying.TimedAttemptSettings;
+import com.google.cloud.storage.StorageRetryStrategy;
+
+import java.util.concurrent.CancellationException;
+
+public class ShouldRetryDecorator<T> implements StorageRetryStrategy {
+
+    private final ResultRetryAlgorithm<T> idempotentRetryAlgorithm;
+    private final ResultRetryAlgorithm<T> nonIdempotentRetryAlgorithm;
+
+    /**
+     * Decorate the should-retry logic for the specified {@link StorageRetryStrategy}
+     *
+     * @param delegate The underling storage retry strategy
+     * @param shouldRetryDecorator The decorated behaviour of {@link ResultRetryAlgorithm#shouldRetry(Throwable, Object)}
+     * @return A decorated {@link StorageRetryStrategy}
+     */
+    public static StorageRetryStrategy decorate(StorageRetryStrategy delegate, Decorator<?> shouldRetryDecorator) {
+        return new ShouldRetryDecorator<>(delegate, shouldRetryDecorator);
+    }
+
+    /**
+     * The logic to use for {@link ResultRetryAlgorithm#shouldRetry(Throwable, Object)}
+     */
+    public interface Decorator<R> {
+        boolean shouldRetry(Throwable prevThrowable, R prevResponse, ResultRetryAlgorithm<R> delegate);
+    }
+
+    /**
+     * @param delegate The delegate {@link StorageRetryStrategy}
+     * @param shouldRetryDecorator The function to call for shouldRetry for idempotent and non-idempotent requests
+     */
+    @SuppressWarnings("unchecked")
+    private ShouldRetryDecorator(StorageRetryStrategy delegate, Decorator<T> shouldRetryDecorator) {
+        this.idempotentRetryAlgorithm = new DelegatingResultRetryAlgorithm<>(
+            (ResultRetryAlgorithm<T>) delegate.getIdempotentHandler(),
+            shouldRetryDecorator
+        );
+        this.nonIdempotentRetryAlgorithm = new DelegatingResultRetryAlgorithm<>(
+            (ResultRetryAlgorithm<T>) delegate.getNonidempotentHandler(),
+            shouldRetryDecorator
+        );
+    }
+
+    @Override
+    public ResultRetryAlgorithm<?> getIdempotentHandler() {
+        return idempotentRetryAlgorithm;
+    }
+
+    @Override
+    public ResultRetryAlgorithm<?> getNonidempotentHandler() {
+        return nonIdempotentRetryAlgorithm;
+    }
+
+    private record DelegatingResultRetryAlgorithm<R>(ResultRetryAlgorithm<R> delegate, Decorator<R> shouldRetryDecorator)
+        implements
+            ResultRetryAlgorithm<R> {
+
+        @Override
+        public TimedAttemptSettings createNextAttempt(Throwable prevThrowable, R prevResponse, TimedAttemptSettings prevSettings) {
+            return delegate.createNextAttempt(prevThrowable, prevResponse, prevSettings);
+        }
+
+        @Override
+        public boolean shouldRetry(Throwable prevThrowable, R prevResponse) throws CancellationException {
+            return shouldRetryDecorator.shouldRetry(prevThrowable, prevResponse, delegate);
+        }
+    }
+}