Add integration test for Azure multi block uploads (#69267)

Relates #68957

plugins/repository-azure/src/internalClusterTest/java/org/elasticsearch/repositories/azure/AzureBlobStoreRepositoryTests.java

@@ -23,6 +23,7 @@ import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.common.settings.MockSecureSettings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeUnit;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.repositories.blobstore.ESMockAPIBasedRepositoryIntegTestCase;
 import org.elasticsearch.rest.RestStatus;
@@ -55,6 +56,7 @@ public class AzureBlobStoreRepositoryTests extends ESMockAPIBasedRepositoryInteg
     protected Settings repositorySettings(String repoName) {
         Settings.Builder settingsBuilder = Settings.builder()
                 .put(super.repositorySettings(repoName))
+                .put(AzureRepository.Repository.MAX_SINGLE_PART_UPLOAD_SIZE_SETTING.getKey(), new ByteSizeValue(1, ByteSizeUnit.MB))
                 .put(AzureRepository.Repository.CONTAINER_SETTING.getKey(), "container")
                 .put(AzureStorageSettings.ACCOUNT_SETTING.getKey(), "test");
         if (randomBoolean()) {
@@ -120,11 +122,6 @@ public class AzureBlobStoreRepositoryTests extends ESMockAPIBasedRepositoryInteg
                 long getUploadBlockSize() {
                     return ByteSizeUnit.MB.toBytes(1);
                 }
-
-                @Override
-                long getSizeThresholdForMultiBlockUpload() {
-                    return ByteSizeUnit.MB.toBytes(1);
-                }
             };
         }
     }

plugins/repository-azure/src/internalClusterTest/java/org/elasticsearch/repositories/azure/AzureStorageCleanupThirdPartyTests.java

@@ -8,27 +8,31 @@
 
 package org.elasticsearch.repositories.azure;
 
-import static org.hamcrest.Matchers.blankOrNullString;
-import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.not;
-
-import java.net.HttpURLConnection;
-import java.util.Collection;
-
+import com.azure.storage.blob.BlobContainerClient;
+import com.azure.storage.blob.BlobServiceClient;
+import com.azure.storage.blob.models.BlobStorageException;
 import org.elasticsearch.action.ActionRunnable;
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.action.support.master.AcknowledgedResponse;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.UUIDs;
+import org.elasticsearch.common.blobstore.BlobContainer;
 import org.elasticsearch.common.settings.MockSecureSettings;
 import org.elasticsearch.common.settings.SecureSettings;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.unit.ByteSizeUnit;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.repositories.AbstractThirdPartyRepositoryTestCase;
 import org.elasticsearch.repositories.blobstore.BlobStoreRepository;
 
-import com.azure.storage.blob.BlobContainerClient;
-import com.azure.storage.blob.BlobServiceClient;
-import com.azure.storage.blob.models.BlobStorageException;
+import java.io.ByteArrayInputStream;
+import java.net.HttpURLConnection;
+import java.util.Collection;
+
+import static org.hamcrest.Matchers.blankOrNullString;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.not;
 
 public class AzureStorageCleanupThirdPartyTests extends AbstractThirdPartyRepositoryTestCase {
 
@@ -78,6 +82,7 @@ public class AzureStorageCleanupThirdPartyTests extends AbstractThirdPartyReposi
             .setSettings(Settings.builder()
                 .put("container", System.getProperty("test.azure.container"))
                 .put("base_path", System.getProperty("test.azure.base"))
+                .put("max_single_part_upload_size", new ByteSizeValue(1, ByteSizeUnit.MB))
             ).get();
         assertThat(putRepositoryResponse.isAcknowledged(), equalTo(true));
         if (Strings.hasText(System.getProperty("test.azure.sas_token"))) {
@@ -111,4 +116,18 @@ public class AzureStorageCleanupThirdPartyTests extends AbstractThirdPartyReposi
         }));
         future.actionGet();
     }
+
+    public void testMultiBlockUpload() throws Exception {
+        final BlobStoreRepository repo = getRepository();
+        // The configured threshold for this test suite is 1mb
+        final int blobSize = ByteSizeUnit.MB.toIntBytes(2);
+        PlainActionFuture<Void> future = PlainActionFuture.newFuture();
+        repo.threadPool().generic().execute(ActionRunnable.run(future, () -> {
+            final BlobContainer blobContainer = repo.blobStore().blobContainer(repo.basePath().add("large_write"));
+            blobContainer.writeBlob(UUIDs.base64UUID(),
+                new ByteArrayInputStream(randomByteArrayOfLength(blobSize)), blobSize, false);
+            blobContainer.delete();
+        }));
+        future.get();
+    }
 }

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

@@ -83,6 +83,7 @@ public class AzureBlobStore implements BlobStore {
     private final String clientName;
     private final String container;
     private final LocationMode locationMode;
+    private final ByteSizeValue maxSinglePartUploadSize;
 
     private final Stats stats = new Stats();
     private final BiConsumer<String, URL> statsConsumer;
@@ -93,6 +94,7 @@ public class AzureBlobStore implements BlobStore {
         this.service = service;
         // locationMode is set per repository, not per client
         this.locationMode = Repository.LOCATION_MODE_SETTING.get(metadata.settings());
+        this.maxSinglePartUploadSize = Repository.MAX_SINGLE_PART_UPLOAD_SIZE_SETTING.get(metadata.settings());
 
         List<RequestStatsCollector> requestStatsCollectors = List.of(
             RequestStatsCollector.create(
@@ -552,7 +554,7 @@ public class AzureBlobStore implements BlobStore {
     }
 
     long getLargeBlobThresholdInBytes() {
-        return service.getSizeThresholdForMultiBlockUpload();
+        return maxSinglePartUploadSize.getBytes();
     }
 
     long getUploadBlockSize() {

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

@@ -18,6 +18,7 @@ import org.elasticsearch.common.blobstore.BlobPath;
 import org.elasticsearch.common.blobstore.BlobStore;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Setting.Property;
+import org.elasticsearch.common.unit.ByteSizeUnit;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.xcontent.NamedXContentRegistry;
@@ -60,6 +61,10 @@ public class AzureRepository extends MeteredBlobStoreRepository {
         public static final Setting<ByteSizeValue> CHUNK_SIZE_SETTING =
             Setting.byteSizeSetting("chunk_size", MAX_CHUNK_SIZE, MIN_CHUNK_SIZE, MAX_CHUNK_SIZE, Property.NodeScope);
         public static final Setting<Boolean> READONLY_SETTING = Setting.boolSetting(READONLY_SETTING_KEY, false, Property.NodeScope);
+        // see ModelHelper.BLOB_DEFAULT_MAX_SINGLE_UPLOAD_SIZE
+        private static final ByteSizeValue DEFAULT_MAX_SINGLE_UPLOAD_SIZE = new ByteSizeValue(256, ByteSizeUnit.MB);
+        public static final Setting<ByteSizeValue> MAX_SINGLE_PART_UPLOAD_SIZE_SETTING =
+            Setting.byteSizeSetting("max_single_part_upload_size", DEFAULT_MAX_SINGLE_UPLOAD_SIZE, Property.NodeScope);
     }
 
     private final ByteSizeValue chunkSize;

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

@@ -66,8 +66,6 @@ public class AzureStorageService {
      */
     public static final ByteSizeValue MAX_CHUNK_SIZE = new ByteSizeValue(MAX_BLOB_SIZE , ByteSizeUnit.BYTES);
 
-    // see ModelHelper.BLOB_DEFAULT_MAX_SINGLE_UPLOAD_SIZE
-    private static final long DEFAULT_MAX_SINGLE_UPLOAD_SIZE = new ByteSizeValue(256, ByteSizeUnit.MB).getBytes();
     private static final long DEFAULT_UPLOAD_BLOCK_SIZE = DEFAULT_BLOCK_SIZE.getBytes();
 
     // 'package' for testing
@@ -123,11 +121,6 @@ public class AzureStorageService {
         return DEFAULT_UPLOAD_BLOCK_SIZE;
     }
 
-    // non-static, package private for testing
-    long getSizeThresholdForMultiBlockUpload() {
-        return DEFAULT_MAX_SINGLE_UPLOAD_SIZE;
-    }
-
     int getMaxReadRetries(String clientName) {
         AzureStorageSettings azureStorageSettings = getClientSettings(clientName);
         return azureStorageSettings.getMaxRetries();

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

@@ -28,6 +28,7 @@ import org.elasticsearch.common.network.InetAddresses;
 import org.elasticsearch.common.settings.MockSecureSettings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeUnit;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.common.util.concurrent.CountDown;
 import org.elasticsearch.mocksocket.MockHttpServer;
@@ -64,6 +65,7 @@ import java.util.stream.Collectors;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.elasticsearch.repositories.azure.AzureRepository.Repository.CONTAINER_SETTING;
 import static org.elasticsearch.repositories.azure.AzureRepository.Repository.LOCATION_MODE_SETTING;
+import static org.elasticsearch.repositories.azure.AzureRepository.Repository.MAX_SINGLE_PART_UPLOAD_SIZE_SETTING;
 import static org.elasticsearch.repositories.azure.AzureStorageSettings.ACCOUNT_SETTING;
 import static org.elasticsearch.repositories.azure.AzureStorageSettings.ENDPOINT_SUFFIX_SETTING;
 import static org.elasticsearch.repositories.azure.AzureStorageSettings.KEY_SETTING;
@@ -156,11 +158,6 @@ public class AzureBlobContainerRetriesTests extends ESTestCase {
                 return ByteSizeUnit.MB.toBytes(1);
             }
 
-            @Override
-            long getSizeThresholdForMultiBlockUpload() {
-                return ByteSizeUnit.MB.toBytes(1);
-            }
-
             @Override
             int getMaxReadRetries(String clientName) {
                 return maxRetries;
@@ -172,6 +169,7 @@ public class AzureBlobContainerRetriesTests extends ESTestCase {
                 .put(CONTAINER_SETTING.getKey(), "container")
                 .put(ACCOUNT_SETTING.getKey(), clientName)
                 .put(LOCATION_MODE_SETTING.getKey(), locationMode)
+                .put(MAX_SINGLE_PART_UPLOAD_SIZE_SETTING.getKey(), new ByteSizeValue(1, ByteSizeUnit.MB))
                 .build());
 
         return new AzureBlobContainer(BlobPath.cleanPath(), new AzureBlobStore(repositoryMetadata, service));