Fix Bug in Azure Repo Exception Handling

plugins/repository-azure/src/main/java/org/elasticsearch/repositories/azure/AzureBlobContainer.java

@@ -62,7 +62,7 @@ private boolean blobExists(String blobName) {
         logger.trace("blobExists({})", blobName);
         try {
             return blobStore.blobExists(buildKey(blobName));
-        } catch (URISyntaxException | StorageException e) {
+        } catch (URISyntaxException | StorageException | IOException e) {
             logger.warn("can not access [{}] in container {{}}: {}", blobName, blobStore, e.getMessage());
         }
         return false;
@@ -97,7 +97,6 @@ public InputStream readBlob(String blobName) throws IOException {
     @Override
     public void writeBlob(String blobName, InputStream inputStream, long blobSize, boolean failIfAlreadyExists) throws IOException {
         logger.trace("writeBlob({}, stream, {})", buildKey(blobName), blobSize);
-
         try {
             blobStore.writeBlob(buildKey(blobName), inputStream, blobSize, failIfAlreadyExists);
         } catch (URISyntaxException|StorageException e) {

plugins/repository-azure/src/main/java/org/elasticsearch/repositories/azure/AzureBlobStore.java

@@ -33,7 +33,6 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URISyntaxException;
-import java.nio.file.FileAlreadyExistsException;
 import java.util.Collections;
 import java.util.Map;
 import java.util.concurrent.Executor;
@@ -88,11 +87,11 @@ public BlobContainer blobContainer(BlobPath path) {
     public void close() {
     }
 
-    public boolean blobExists(String blob) throws URISyntaxException, StorageException {
+    public boolean blobExists(String blob) throws URISyntaxException, StorageException, IOException {
         return service.blobExists(clientName, container, blob);
     }
 
-    public void deleteBlob(String blob) throws URISyntaxException, StorageException {
+    public void deleteBlob(String blob) throws URISyntaxException, StorageException, IOException {
         service.deleteBlob(clientName, container, blob);
     }
 
@@ -106,17 +105,17 @@ public InputStream getInputStream(String blob) throws URISyntaxException, Storag
     }
 
     public Map<String, BlobMetaData> listBlobsByPrefix(String keyPath, String prefix)
-        throws URISyntaxException, StorageException {
+        throws URISyntaxException, StorageException, IOException {
         return service.listBlobsByPrefix(clientName, container, keyPath, prefix);
     }
 
-    public Map<String, BlobContainer> children(BlobPath path) throws URISyntaxException, StorageException {
+    public Map<String, BlobContainer> children(BlobPath path) throws URISyntaxException, StorageException, IOException {
         return Collections.unmodifiableMap(service.children(clientName, container, path).stream().collect(
             Collectors.toMap(Function.identity(), name -> new AzureBlobContainer(path.add(name), this, threadPool))));
     }
 
     public void writeBlob(String blobName, InputStream inputStream, long blobSize, boolean failIfAlreadyExists)
-        throws URISyntaxException, StorageException, FileAlreadyExistsException {
+        throws URISyntaxException, StorageException, IOException {
         service.writeBlob(this.clientName, container, blobName, inputStream, blobSize, failIfAlreadyExists);
     }
 }

plugins/repository-azure/src/main/java/org/elasticsearch/repositories/azure/AzureStorageService.java

@@ -177,7 +177,7 @@ static String blobNameFromUri(URI uri) {
         return splits[2];
     }
 
-    public boolean blobExists(String account, String container, String blob) throws URISyntaxException, StorageException {
+    public boolean blobExists(String account, String container, String blob) throws URISyntaxException, StorageException, IOException {
         // Container name must be lower case.
         final Tuple<CloudBlobClient, Supplier<OperationContext>> client = client(account);
         final CloudBlobContainer blobContainer = client.v1().getContainerReference(container);
@@ -187,7 +187,7 @@ public boolean blobExists(String account, String container, String blob) throws
         });
     }
 
-    public void deleteBlob(String account, String container, String blob) throws URISyntaxException, StorageException {
+    public void deleteBlob(String account, String container, String blob) throws URISyntaxException, StorageException, IOException {
         final Tuple<CloudBlobClient, Supplier<OperationContext>> client = client(account);
         // Container name must be lower case.
         final CloudBlobContainer blobContainer = client.v1().getContainerReference(container);
@@ -267,7 +267,7 @@ public InputStream getInputStream(String account, String container, String blob)
     }
 
     public Map<String, BlobMetaData> listBlobsByPrefix(String account, String container, String keyPath, String prefix)
-        throws URISyntaxException, StorageException {
+            throws URISyntaxException, StorageException, IOException {
         // NOTE: this should be here: if (prefix == null) prefix = "";
         // however, this is really inefficient since deleteBlobsByPrefix enumerates everything and
         // then does a prefix match on the result; it should just call listBlobsByPrefix with the prefix!
@@ -295,7 +295,7 @@ public Map<String, BlobMetaData> listBlobsByPrefix(String account, String contai
         return Map.copyOf(blobsBuilder);
     }
 
-    public Set<String> children(String account, String container, BlobPath path) throws URISyntaxException, StorageException {
+    public Set<String> children(String account, String container, BlobPath path) throws URISyntaxException, StorageException, IOException {
         final var blobsBuilder = new HashSet<String>();
         final Tuple<CloudBlobClient, Supplier<OperationContext>> client = client(account);
         final CloudBlobContainer blobContainer = client.v1().getContainerReference(container);
@@ -319,8 +319,9 @@ public Set<String> children(String account, String container, BlobPath path) thr
     }
 
     public void writeBlob(String account, String container, String blobName, InputStream inputStream, long blobSize,
-                          boolean failIfAlreadyExists)
-        throws URISyntaxException, StorageException, FileAlreadyExistsException {
+                          boolean failIfAlreadyExists) throws URISyntaxException, StorageException, IOException {
+        assert inputStream.markSupported()
+            : "Should not be used with non-mark supporting streams as their retry handling in the SDK is broken";
         logger.trace(() -> new ParameterizedMessage("writeBlob({}, stream, {})", blobName, blobSize));
         final Tuple<CloudBlobClient, Supplier<OperationContext>> client = client(account);
         final CloudBlobContainer blobContainer = client.v1().getContainerReference(container);

plugins/repository-azure/src/main/java/org/elasticsearch/repositories/azure/SocketAccess.java

@@ -44,33 +44,51 @@ public static <T> T doPrivilegedIOException(PrivilegedExceptionAction<T> operati
         try {
             return AccessController.doPrivileged(operation);
         } catch (PrivilegedActionException e) {
-            throw (IOException) e.getCause();
+            Throwable cause = e.getCause();
+            if (cause instanceof IOException) {
+                throw (IOException) e.getCause();
+            } else {
+                throw new IOException(cause);
+            }
         }
     }
 
-    public static <T> T doPrivilegedException(PrivilegedExceptionAction<T> operation) throws StorageException {
+    public static <T> T doPrivilegedException(PrivilegedExceptionAction<T> operation)
+            throws StorageException, IOException, URISyntaxException {
         SpecialPermission.check();
         try {
             return AccessController.doPrivileged(operation);
         } catch (PrivilegedActionException e) {
-            throw (StorageException) e.getCause();
+            safeRethrowCause(e);
+            assert false : "always throws";
+            return null;
         }
     }
 
-    public static void doPrivilegedVoidException(StorageRunnable action) throws StorageException, URISyntaxException {
+    public static void doPrivilegedVoidException(StorageRunnable action) throws StorageException, URISyntaxException, IOException {
         SpecialPermission.check();
         try {
             AccessController.doPrivileged((PrivilegedExceptionAction<Void>) () -> {
                 action.executeCouldThrow();
                 return null;
             });
         } catch (PrivilegedActionException e) {
-            Throwable cause = e.getCause();
-            if (cause instanceof StorageException) {
-                throw (StorageException) cause;
-            } else {
-                throw (URISyntaxException) cause;
-            }
+            safeRethrowCause(e);
+        }
+    }
+
+    private static void safeRethrowCause(PrivilegedActionException e) throws IOException, URISyntaxException, StorageException {
+        Throwable cause = e.getCause();
+        if (cause instanceof StorageException) {
+            throw (StorageException) cause;
+        } else if (cause instanceof URISyntaxException) {
+            throw (URISyntaxException) cause;
+        } else if (cause instanceof IOException) {
+            throw (IOException) cause;
+        } else if (cause instanceof RuntimeException) {
+            throw (RuntimeException) cause;
+        } else {
+            throw new IOException(cause);
         }
     }
 

plugins/repository-azure/src/test/java/org/elasticsearch/repositories/azure/AzureBlobContainerRetriesTests.java

@@ -49,6 +49,7 @@
 import org.junit.Before;
 
 import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.net.InetAddress;
@@ -63,6 +64,7 @@
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -294,6 +296,44 @@ public void testWriteLargeBlob() throws Exception {
         assertThat(blocks.isEmpty(), is(true));
     }
 
+    public void testRetryUntilFail() throws IOException {
+        final AtomicBoolean requestReceived = new AtomicBoolean(false);
+        httpServer.createContext("/container/write_blob_max_retries", exchange -> {
+            try {
+                if (requestReceived.compareAndSet(false, true)) {
+                    throw new AssertionError("Should not receive two requests");
+                } else {
+                    exchange.sendResponseHeaders(RestStatus.CREATED.getStatus(), -1);
+                }
+            } finally {
+                exchange.close();
+            }
+        });
+
+        final BlobContainer blobContainer = createBlobContainer(randomIntBetween(2, 5));
+        try (InputStream stream = new InputStream() {
+
+            @Override
+            public int read() throws IOException {
+                throw new IOException("foo");
+            }
+
+            @Override
+            public boolean markSupported() {
+                return true;
+            }
+
+            @Override
+            public void reset() {
+                throw new AssertionError("should not be called");
+            }
+        }) {
+            final IOException ioe = expectThrows(IOException.class, () ->
+                blobContainer.writeBlob("write_blob_max_retries", stream, randomIntBetween(1, 128), randomBoolean()));
+            assertThat(ioe.getMessage(), is("foo"));
+        }
+    }
+
     private static byte[] randomBlobContent() {
         return randomByteArrayOfLength(randomIntBetween(1, frequently() ? 512 : 1 << 20)); // rarely up to 1mb
     }