Migrate GCS repository plugin to use GCS SDK v 2.x

willyborankin · willyborankin · commit ad267244f55f · 2025-08-19T21:13:35.000+02:00
Signed-off-by: Andrey Pleskach &lt;ples@aiven.io&gt;
diff --git a/buildSrc/src/main/groovy/org/opensearch/gradle/test/TestWithSslPlugin.java b/buildSrc/src/main/groovy/org/opensearch/gradle/test/TestWithSslPlugin.java
@@ -77,6 +77,7 @@ public void apply(Project project) {
             File keystoreDir = new File(project.getBuildDir(), "keystore/test/ssl");
             File nodeKeystore = new File(keystoreDir, "test-node.jks");
             File clientKeyStore = new File(keystoreDir, "test-client.jks");
+            @SuppressWarnings("unchecked")
             NamedDomainObjectContainer<OpenSearchCluster> clusters = (NamedDomainObjectContainer<OpenSearchCluster>) project.getExtensions()
                 .getByName(TestClustersPlugin.EXTENSION_NAME);
             clusters.all(c -> {
diff --git a/plugins/repository-gcs/src/main/java/org/opensearch/repositories/gcs/GoogleCloudStorageBlobStore.java b/plugins/repository-gcs/src/main/java/org/opensearch/repositories/gcs/GoogleCloudStorageBlobStore.java
@@ -300,7 +300,18 @@ private void writeBlobResumable(BlobInfo blobInfo, InputStream inputStream, long
                     @SuppressForbidden(reason = "channel is based on a socket")
                     @Override
                     public int write(final ByteBuffer src) throws IOException {
-                        return SocketAccess.doPrivilegedIOException(() -> writeChannel.write(src));
+                        try {
+                            return SocketAccess.doPrivilegedIOException(() -> writeChannel.write(src));
+                        } catch (final IOException ioe) {
+                            final StorageException storageException = (StorageException) ExceptionsHelper.unwrap(
+                                ioe,
+                                StorageException.class
+                            );
+                            if (storageException != null) {
+                                throw storageException;
+                            }
+                            throw ioe;
+                        }
                     }
 
                     @Override
diff --git a/plugins/repository-gcs/src/main/java/org/opensearch/repositories/gcs/GoogleCloudStorageService.java b/plugins/repository-gcs/src/main/java/org/opensearch/repositories/gcs/GoogleCloudStorageService.java
@@ -216,7 +216,8 @@ StorageOptions createStorageOptions(
                     mapBuilder.put("user-agent", clientSettings.getApplicationName());
                 }
                 return mapBuilder.immutableMap();
-            });
+            })
+            .setStorageRetryStrategy(new GoogleShouldRetryStorageStrategy());
         if (Strings.hasLength(clientSettings.getHost())) {
             storageOptionsBuilder.setHost(clientSettings.getHost());
         }
diff --git a/plugins/repository-gcs/src/main/java/org/opensearch/repositories/gcs/GoogleShouldRetryStorageStrategy.java b/plugins/repository-gcs/src/main/java/org/opensearch/repositories/gcs/GoogleShouldRetryStorageStrategy.java
@@ -0,0 +1,63 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.repositories.gcs;
+
+import com.google.api.gax.retrying.ResultRetryAlgorithm;
+import com.google.api.gax.retrying.TimedAttemptSettings;
+import com.google.cloud.BaseService;
+import com.google.cloud.storage.StorageRetryStrategy;
+import org.opensearch.ExceptionsHelper;
+
+import java.net.SocketException;
+import java.net.UnknownHostException;
+import java.util.concurrent.CancellationException;
+
+import static java.util.Objects.nonNull;
+
+public class GoogleShouldRetryStorageStrategy implements StorageRetryStrategy {
+
+    private final DelagateResultRetryAlgorithm<?> idempotentHandler = new DelagateResultRetryAlgorithm<>(BaseService.EXCEPTION_HANDLER);
+
+    private final DelagateResultRetryAlgorithm<?> nonIdempotentHandler = new DelagateResultRetryAlgorithm<>(BaseService.EXCEPTION_HANDLER);
+
+    private static final class DelagateResultRetryAlgorithm<T> implements ResultRetryAlgorithm<T> {
+
+        private final ResultRetryAlgorithm<T> resultRetryAlgorithm;
+
+        private DelagateResultRetryAlgorithm(ResultRetryAlgorithm<T> resultRetryAlgorithm) {
+            this.resultRetryAlgorithm = resultRetryAlgorithm;
+        }
+
+        @Override
+        public TimedAttemptSettings createNextAttempt(Throwable prevThrowable, T prevResponse, TimedAttemptSettings prevSettings) {
+            return resultRetryAlgorithm.createNextAttempt(prevThrowable, prevResponse, prevSettings);
+        }
+
+        @Override
+        public boolean shouldRetry(Throwable prevThrowable, T prevResponse) throws CancellationException {
+            if (nonNull(ExceptionsHelper.unwrap(prevThrowable, UnknownHostException.class))) {
+                return true;
+            }
+            if (nonNull(ExceptionsHelper.unwrap(prevThrowable, SocketException.class))) {
+                return true;
+            }
+            return resultRetryAlgorithm.shouldRetry(prevThrowable, prevResponse);
+        }
+    };
+
+    @Override
+    public ResultRetryAlgorithm<?> getIdempotentHandler() {
+        return idempotentHandler;
+    }
+
+    @Override
+    public ResultRetryAlgorithm<?> getNonidempotentHandler() {
+        return nonIdempotentHandler;
+    }
+}
diff --git a/plugins/repository-gcs/src/test/java/org/opensearch/repositories/gcs/GoogleCloudStorageBlobContainerRetriesTests.java b/plugins/repository-gcs/src/test/java/org/opensearch/repositories/gcs/GoogleCloudStorageBlobContainerRetriesTests.java
@@ -52,7 +52,6 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.common.util.concurrent.CountDown;
-import org.opensearch.core.common.Strings;
 import org.opensearch.core.common.bytes.BytesArray;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.common.unit.ByteSizeValue;
@@ -75,6 +74,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
+import fixture.gcs.ContentHttpHeadersParser;
 import fixture.gcs.FakeOAuth2HttpHandler;
 import org.threeten.bp.Duration;
 
@@ -92,9 +92,6 @@
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.notNullValue;
-import static fixture.gcs.GoogleCloudStorageHttpHandler.getContentRangeEnd;
-import static fixture.gcs.GoogleCloudStorageHttpHandler.getContentRangeLimit;
-import static fixture.gcs.GoogleCloudStorageHttpHandler.getContentRangeStart;
 import static fixture.gcs.GoogleCloudStorageHttpHandler.parseMultipartRequestBody;
 
 @SuppressForbidden(reason = "use a http server")
@@ -152,7 +149,6 @@ StorageOptions createStorageOptions(
                     .setInitialRetryDelay(Duration.ofMillis(10L))
                     .setRetryDelayMultiplier(1.0d)
                     .setMaxRetryDelay(Duration.ofSeconds(1L))
-                    .setJittered(false)
                     .setInitialRpcTimeout(Duration.ofSeconds(1))
                     .setRpcTimeoutMultiplier(options.getRetrySettings().getRpcTimeoutMultiplier())
                     .setMaxRpcTimeout(Duration.ofSeconds(1));
@@ -163,6 +159,7 @@ StorageOptions createStorageOptions(
                     .setHost(options.getHost())
                     .setCredentials(options.getCredentials())
                     .setRetrySettings(retrySettingsBuilder.build())
+                    .setStorageRetryStrategy(new GoogleShouldRetryStorageStrategy())
                     .build();
             }
         };
@@ -220,7 +217,8 @@ public void testWriteBlobWithRetries() throws Exception {
                 assertThat(content.isPresent(), is(true));
                 assertThat(content.get().v1(), equalTo("write_blob_max_retries"));
                 if (Objects.deepEquals(bytes, BytesReference.toBytes(content.get().v2()))) {
-                    byte[] response = ("{\"bucket\":\"bucket\",\"name\":\"" + content.get().v1() + "\"}").getBytes(UTF_8);
+                    byte[] response = String.format("""
+                        {"bucket":"bucket","name":"%s"}""", content.get().v1()).getBytes(UTF_8);
                     exchange.getResponseHeaders().add("Content-Type", "application/json");
                     exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
                     exchange.getResponseBody().write(response);
@@ -274,8 +272,7 @@ public void testWriteBlobWithReadTimeouts() {
     }
 
     public void testWriteLargeBlob() throws IOException {
-        // See {@link BaseWriteChannel#DEFAULT_CHUNK_SIZE}
-        final int defaultChunkSize = 60 * 256 * 1024;
+        final int defaultChunkSize = Math.toIntExact(ByteSizeValue.parseBytesSizeValue("16mb", "aaa").getBytes());
         final int nbChunks = randomIntBetween(3, 5);
         final int lastChunkSize = randomIntBetween(1, defaultChunkSize - 1);
         final int totalChunks = nbChunks + 1;
@@ -295,6 +292,7 @@ public void testWriteLargeBlob() throws IOException {
         final AtomicInteger countUploads = new AtomicInteger(nbErrors * totalChunks);
         final AtomicBoolean allow410Gone = new AtomicBoolean(randomBoolean());
         final AtomicBoolean allowReadTimeout = new AtomicBoolean(rarely());
+        final AtomicInteger bytesReceived = new AtomicInteger();
         final int wrongChunk = randomIntBetween(1, totalChunks);
 
         final AtomicReference<String> sessionUploadId = new AtomicReference<>(UUIDs.randomBase64UUID());
@@ -325,7 +323,6 @@ public void testWriteLargeBlob() throws IOException {
                     assertThat(wrongChunk, greaterThan(0));
                     return;
                 }
-
             } else if ("PUT".equals(exchange.getRequestMethod())) {
                 final String uploadId = params.get("upload_id");
                 if (uploadId.equals(sessionUploadId.get()) == false) {
@@ -348,29 +345,43 @@ public void testWriteLargeBlob() throws IOException {
                         // we must reset the counters because the whole object upload will be retried
                         countInits.set(nbErrors);
                         countUploads.set(nbErrors * totalChunks);
+                        bytesReceived.set(0);
 
                         exchange.sendResponseHeaders(HttpStatus.SC_GONE, -1);
                         return;
                     }
                 }
 
-                final String range = exchange.getRequestHeaders().getFirst("Content-Range");
-                assertTrue(Strings.hasLength(range));
+                final String contentRangeHeaderValue = exchange.getRequestHeaders().getFirst("Content-Range");
+                final var contentRange = ContentHttpHeadersParser.parseContentRangeHeader(contentRangeHeaderValue);
+                assertNotNull("Invalid content range header: " + contentRangeHeaderValue, contentRange);
+
+                if (!contentRange.hasRange()) {
+                    // Content-Range: */... is a status check
+                    // https://cloud.google.com/storage/docs/performing-resumable-uploads#status-check
+                    final int receivedSoFar = bytesReceived.get();
+                    if (receivedSoFar > 0) {
+                        exchange.getResponseHeaders().add("Range", String.format("bytes=0-%s", receivedSoFar));
+                    }
+                    exchange.getResponseHeaders().add("Content-Length", "0");
+                    exchange.sendResponseHeaders(308 /* Resume Incomplete */, -1);
+                    return;
+                }
 
                 if (countUploads.decrementAndGet() % 2 == 0) {
                     assertThat(Math.toIntExact(requestBody.length()), anyOf(equalTo(defaultChunkSize), equalTo(lastChunkSize)));
-
-                    final int rangeStart = getContentRangeStart(range);
-                    final int rangeEnd = getContentRangeEnd(range);
+                    final int rangeStart = Math.toIntExact(contentRange.start());
+                    final int rangeEnd = Math.toIntExact(contentRange.end());
                     assertThat(rangeEnd + 1 - rangeStart, equalTo(Math.toIntExact(requestBody.length())));
                     assertThat(new BytesArray(data, rangeStart, rangeEnd - rangeStart + 1), is(requestBody));
+                    bytesReceived.updateAndGet(existing -> Math.max(existing, rangeEnd));
 
-                    final Integer limit = getContentRangeLimit(range);
-                    if (limit != null) {
+                    if (contentRange.size() != null) {
+                        exchange.getResponseHeaders().add("x-goog-stored-content-length", String.valueOf(bytesReceived.get() + 1));
                         exchange.sendResponseHeaders(RestStatus.OK.getStatus(), -1);
                         return;
                     } else {
-                        exchange.getResponseHeaders().add("Range", String.format(Locale.ROOT, "bytes=%d/%d", rangeStart, rangeEnd));
+                        exchange.getResponseHeaders().add("Range", String.format("bytes=%s-%s", rangeStart, rangeEnd));
                         exchange.getResponseHeaders().add("Content-Length", "0");
                         exchange.sendResponseHeaders(308 /* Resume Incomplete */, -1);
                         return;
diff --git a/test/fixtures/gcs-fixture/src/main/java/fixture/gcs/ContentHttpHeadersParser.java b/test/fixtures/gcs-fixture/src/main/java/fixture/gcs/ContentHttpHeadersParser.java
@@ -0,0 +1,82 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package fixture.gcs;
+
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+public class ContentHttpHeadersParser {
+
+    private static final Pattern PATTERN_CONTENT_RANGE = Pattern.compile("bytes=([0-9]+)-([0-9]+)");
+    private static final Pattern PATTERN_CONTENT_RANGE_BYTES = Pattern.compile("bytes (?:(\\d+)-(\\d+)|\\*)/(?:(\\d+)|\\*)");
+
+
+    public static Range parseRangeHeader(String rangeHeaderValue) {
+        final Matcher matcher = PATTERN_CONTENT_RANGE.matcher(rangeHeaderValue);
+        if (matcher.matches()) {
+            try {
+                return new Range(Integer.parseInt(matcher.group(1)), Integer.parseInt(matcher.group(2)));
+            } catch (NumberFormatException e) {
+                return null;
+            }
+        }
+        return null;
+    }
+
+    public record Range(int start, int end) {
+
+        public String header() {
+            return String.format("bytes=%s-%s", start, end);
+        }
+
+    }
+
+    public static ContentRange parseContentRangeHeader(String contentRangeHeaderValue) {
+        final Matcher matcher = PATTERN_CONTENT_RANGE_BYTES.matcher(contentRangeHeaderValue);
+        if (matcher.matches()) {
+            try {
+                if (matcher.groupCount() == 3) {
+                    final Integer start = parseIntegerValue(matcher.group(1));
+                    final Integer end = parseIntegerValue(matcher.group(2));
+                    final Integer size = parseIntegerValue(matcher.group(3));
+                    return new ContentRange(start, end, size);
+                }
+            } catch (NumberFormatException e) {
+                return null;
+            }
+        }
+        return null;
+    }
+
+    private static Integer parseIntegerValue(String value) {
+        return value == null ? null : Integer.parseInt(value);
+    }
+
+    public record ContentRange(Integer start, Integer end, Integer size) {
+
+        public ContentRange {
+            assert (start == null) == (end == null) : "Must have either start and end or neither";
+        }
+
+        public boolean hasRange() {
+            return start != null && end != null;
+        }
+
+        public boolean hasSize() {
+            return size != null;
+        }
+
+        public String headerString() {
+            final String rangeString = hasRange() ? start + "-" + end : "*";
+            final String sizeString = hasSize() ? String.valueOf(size) : "*";
+            return "bytes " + rangeString + "/" + sizeString;
+        }
+    }
+
+}
diff --git a/test/fixtures/gcs-fixture/src/main/java/fixture/gcs/FakeOAuth2HttpHandler.java b/test/fixtures/gcs-fixture/src/main/java/fixture/gcs/FakeOAuth2HttpHandler.java
@@ -48,9 +48,12 @@ public class FakeOAuth2HttpHandler implements HttpHandler {
     @Override
     public void handle(final HttpExchange exchange) throws IOException {
         try {
-            while (exchange.getRequestBody().read(BUFFER) >= 0) ;
-            byte[] response = ("{\"access_token\":\"foo\",\"token_type\":\"Bearer\",\"expires_in\":3600}").getBytes(UTF_8);
+            while (exchange.getRequestBody().read(BUFFER) >= 0) {
+            }
+            byte[] response = """
+                {"access_token":"foo","token_type":"Bearer","expires_in":3600}""".getBytes(UTF_8);
             exchange.getResponseHeaders().add("Content-Type", "application/json");
+            exchange.getResponseHeaders().add("Metadata-Flavor", "Google");
             exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
             exchange.getResponseBody().write(response);
         } finally {
diff --git a/test/fixtures/gcs-fixture/src/main/java/fixture/gcs/GoogleCloudStorageHttpHandler.java b/test/fixtures/gcs-fixture/src/main/java/fixture/gcs/GoogleCloudStorageHttpHandler.java

Original file line number	Diff line number	Diff line change
`@@ -216,7 +216,8 @@ StorageOptions createStorageOptions(`
`216`	`216`	`mapBuilder.put("user-agent", clientSettings.getApplicationName());`
`217`	`217`	`}`
`218`	`218`	`return mapBuilder.immutableMap();`
`219`		`- });`
	`219`	`+ })`
	`220`	`+ .setStorageRetryStrategy(new GoogleShouldRetryStorageStrategy());`
`220`	`221`	`if (Strings.hasLength(clientSettings.getHost())) {`
`221`	`222`	`storageOptionsBuilder.setHost(clientSettings.getHost());`
`222`	`223`	`}`