From 70ee92384428e60e14f3752074203fa836cf82b6 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Mon, 28 Apr 2025 05:16:12 +0530 Subject: [PATCH 01/20] git commit -m "UTs for : executeMultipartUploadIfEtagMatches" --- .../s3/S3BlobStoreContainerTests.java | 1168 +++++++++++++++++ 1 file changed, 1168 insertions(+) diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java index 9e931c717bdf4..dba8044c19c80 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java @@ -64,6 +64,7 @@ import software.amazon.awssdk.services.s3.model.PutObjectRequest; import software.amazon.awssdk.services.s3.model.PutObjectResponse; import software.amazon.awssdk.services.s3.model.S3Error; +import software.amazon.awssdk.services.s3.model.S3Exception; import software.amazon.awssdk.services.s3.model.S3Object; import software.amazon.awssdk.services.s3.model.ServerSideEncryption; import software.amazon.awssdk.services.s3.model.StorageClass; @@ -72,6 +73,7 @@ import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Iterable; import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Publisher; +import org.opensearch.OpenSearchException; import org.opensearch.action.LatchedActionListener; import org.opensearch.common.blobstore.BlobContainer; import org.opensearch.common.blobstore.BlobMetadata; @@ -97,6 +99,7 @@ import java.util.List; import java.util.Map; import java.util.NoSuchElementException; +import java.util.Random; import java.util.Set; import java.util.UUID; import java.util.concurrent.CompletableFuture; @@ -120,6 +123,8 @@ import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.reset; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; @@ -939,6 +944,1169 @@ public void testNumberOfMultiparts() { assertNumberOfMultiparts(factor + 1, remaining, (size * factor) + remaining, size); } + /** + * Tests the basic success path for conditional multipart upload when ETag matches + */ + public void testExecuteMultipartUploadIfEtagMatchesSuccess() throws IOException { + // Setup test parameters with randomized values + final String bucketName = randomAlphaOfLengthBetween(1, 10); + final String blobName = randomAlphaOfLengthBetween(1, 10); + final String inputETag = randomAlphaOfLengthBetween(8, 32); // Input ETag for conditional check + final String finalETag = randomAlphaOfLengthBetween(8, 32); // ETag returned by complete operation + final String uploadId = randomAlphaOfLengthBetween(10, 20); + + // Add metadata - ADDED + final Map metadata = new HashMap<>(); + metadata.put("key1", "value1"); + metadata.put("key2", "value2"); + + // Setup path + final BlobPath blobPath = new BlobPath(); + if (randomBoolean()) { + IntStream.of(randomIntBetween(1, 5)).forEach(value -> blobPath.add("path_" + value)); + } + + // Calculate sizes for multipart upload - ensure minimum viable multipart size + final long partSize = ByteSizeUnit.MB.toBytes(5); // 5MB minimum part size + final int partCount = randomIntBetween(2, 5); + final long lastPartSize = randomIntBetween(1, (int) partSize); + final long blobSize = partSize * (partCount - 1) + lastPartSize; + + // Mock S3BlobStore with randomized parameters - ADDED + final S3BlobStore blobStore = mock(S3BlobStore.class); + final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); + when(blobStore.bucket()).thenReturn(bucketName); + when(blobStore.bufferSizeInBytes()).thenReturn(partSize); + when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); + + // Randomize storage class - ADDED + final StorageClass storageClass = randomFrom(StorageClass.values()); + when(blobStore.getStorageClass()).thenReturn(storageClass); + + // Randomize server-side encryption - ADDED + final boolean serverSideEncryption = randomBoolean(); + when(blobStore.serverSideEncryption()).thenReturn(serverSideEncryption); + + // Randomize ACL - ADDED + final ObjectCannedACL cannedAccessControlList = randomBoolean() ? randomFrom(ObjectCannedACL.values()) : null; + if (cannedAccessControlList != null) { + when(blobStore.getCannedACL()).thenReturn(cannedAccessControlList); + } + + final boolean isUploadRetryEnabled = randomBoolean(); + when(blobStore.isUploadRetryEnabled()).thenReturn(isUploadRetryEnabled); + + final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); + + // Mock S3 client with argument captors to verify requests + final S3Client client = mock(S3Client.class); + final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); + when(blobStore.clientReference()).thenReturn(clientReference); + when(clientReference.get()).thenReturn(client); + + // Setup argument captors for request validation + final ArgumentCaptor createRequestCaptor = ArgumentCaptor.forClass( + CreateMultipartUploadRequest.class + ); + final ArgumentCaptor uploadPartRequestCaptor = ArgumentCaptor.forClass(UploadPartRequest.class); + final ArgumentCaptor requestBodyCaptor = ArgumentCaptor.forClass(RequestBody.class); + final ArgumentCaptor completeRequestCaptor = ArgumentCaptor.forClass( + CompleteMultipartUploadRequest.class + ); + + // Mock createMultipartUpload response with uploadId + when(client.createMultipartUpload(createRequestCaptor.capture())).thenReturn( + CreateMultipartUploadResponse.builder().uploadId(uploadId).build() + ); + + // Generate unique ETags for each part + final List partETags = new ArrayList<>(); + for (int i = 0; i < partCount; i++) { + partETags.add("etag-part-" + (i + 1)); + } + + // Mock uploadPart responses with the appropriate ETags + when(client.uploadPart(uploadPartRequestCaptor.capture(), requestBodyCaptor.capture())).thenAnswer(invocation -> { + UploadPartRequest request = (UploadPartRequest) invocation.getArguments()[0]; + int partNumber = request.partNumber(); + return UploadPartResponse.builder().eTag(partETags.get(partNumber - 1)).build(); + }); + + // Mock successful completion response WITH final ETag - FIXED + when(client.completeMultipartUpload(completeRequestCaptor.capture())).thenReturn( + CompleteMultipartUploadResponse.builder().eTag(finalETag).build() + ); + + // Create action listener to verify callback + @SuppressWarnings("unchecked") + ActionListener etagListener = mock(ActionListener.class); + + // Execute the multipart upload + final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); + blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, metadata, inputETag, etagListener); + + // Verify the create multipart request parameters + final CreateMultipartUploadRequest createRequest = createRequestCaptor.getValue(); + assertEquals(bucketName, createRequest.bucket()); + assertEquals(blobPath.buildAsString() + blobName, createRequest.key()); + assertEquals(storageClass, createRequest.storageClass()); + assertEquals(cannedAccessControlList, createRequest.acl()); + assertEquals(metadata, createRequest.metadata()); + + // Verify server-side encryption if enabled - ADDED + if (serverSideEncryption) { + assertEquals(ServerSideEncryption.AES256, createRequest.serverSideEncryption()); + } + + // Verify part upload requests + List partRequests = uploadPartRequestCaptor.getAllValues(); + assertEquals(partCount, partRequests.size()); + + // Verify request bodies for content (partial content verification) - ADDED + List requestBodies = requestBodyCaptor.getAllValues(); + assertEquals(partCount, requestBodies.size()); + + // Check each part request + for (int i = 0; i < partCount; i++) { + UploadPartRequest partRequest = partRequests.get(i); + assertEquals(bucketName, partRequest.bucket()); + assertEquals(blobPath.buildAsString() + blobName, partRequest.key()); + assertEquals(uploadId, partRequest.uploadId()); + assertEquals(Integer.valueOf(i + 1), partRequest.partNumber()); + // Verify part size - last part can be different + long expectedPartSize = (i < partCount - 1) ? partSize : lastPartSize; + assertEquals(expectedPartSize, partRequest.contentLength().longValue()); + + // Verify content for first part only (partial content verification) - MODIFIED + if (i == 0) { + RequestBody body = requestBodies.get(i); + try (InputStream is = body.contentStreamProvider().newStream()) { + // Just verify the stream exists + assertNotNull(is); + } + } + } + + // Verify complete request + CompleteMultipartUploadRequest completeRequest = completeRequestCaptor.getValue(); + assertEquals(bucketName, completeRequest.bucket()); + assertEquals(blobPath.buildAsString() + blobName, completeRequest.key()); + assertEquals(uploadId, completeRequest.uploadId()); + + // Verify ifMatch on complete request (NOT create request) - FIXED + assertEquals(inputETag, completeRequest.ifMatch()); + + // Verify listener was called with the FINAL ETag (not last part's ETag) - FIXED + verify(etagListener).onResponse(finalETag); + verify(etagListener, never()).onFailure(any()); + + // Verify resources were properly closed + verify(clientReference).close(); + } + + /** + * Tests conditional multipart upload with various configuration combinations + */ + public void testExecuteMultipartUploadIfEtagMatchesWithMetadataAndSSE() throws IOException { + // Setup test parameters + final String bucketName = randomAlphaOfLengthBetween(1, 10); + final String blobName = randomAlphaOfLengthBetween(1, 10); + final String inputETag = randomAlphaOfLengthBetween(8, 32); + final String finalETag = randomAlphaOfLengthBetween(8, 32); // Added distinct final ETag + final String uploadId = randomAlphaOfLengthBetween(10, 20); + + // Create metadata map + final Map metadata = new HashMap<>(); + metadata.put("key1", "value1"); + metadata.put("key2", "value2"); + + final BlobPath blobPath = new BlobPath(); + + // Setup for a minimal viable multipart upload + final long partSize = ByteSizeUnit.MB.toBytes(5); + final long blobSize = partSize * 2; // Two parts + + // Mock S3BlobStore with various configurations enabled + final S3BlobStore blobStore = mock(S3BlobStore.class); + final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); + when(blobStore.bucket()).thenReturn(bucketName); + when(blobStore.bufferSizeInBytes()).thenReturn(partSize); + when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); + + // Setup random configurations + final boolean serverSideEncryption = true; // Force enabled to test this case + when(blobStore.serverSideEncryption()).thenReturn(serverSideEncryption); + + final StorageClass storageClass = randomFrom(StorageClass.values()); + when(blobStore.getStorageClass()).thenReturn(storageClass); + + final ObjectCannedACL cannedAccessControlList = randomFrom(ObjectCannedACL.values()); + when(blobStore.getCannedACL()).thenReturn(cannedAccessControlList); + + final boolean isUploadRetryEnabled = randomBoolean(); + when(blobStore.isUploadRetryEnabled()).thenReturn(isUploadRetryEnabled); + + final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); + + // Mock S3 client and capture requests + final S3Client client = mock(S3Client.class); + final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); + when(blobStore.clientReference()).thenReturn(clientReference); + when(clientReference.get()).thenReturn(client); + + // Setup argument captors + final ArgumentCaptor createRequestCaptor = ArgumentCaptor.forClass( + CreateMultipartUploadRequest.class + ); + final ArgumentCaptor uploadPartRequestCaptor = ArgumentCaptor.forClass(UploadPartRequest.class); + final ArgumentCaptor requestBodyCaptor = ArgumentCaptor.forClass(RequestBody.class); + final ArgumentCaptor completeRequestCaptor = ArgumentCaptor.forClass( + CompleteMultipartUploadRequest.class + ); + + // Mock responses + when(client.createMultipartUpload(createRequestCaptor.capture())).thenReturn( + CreateMultipartUploadResponse.builder().uploadId(uploadId).build() + ); + + // Generate part ETags + final List partETags = List.of("etag-part-1", "etag-part-2"); + + // Mock part upload responses + when(client.uploadPart(uploadPartRequestCaptor.capture(), requestBodyCaptor.capture())).thenAnswer(invocation -> { + UploadPartRequest request = (UploadPartRequest) invocation.getArguments()[0]; + int partNumber = request.partNumber(); + return UploadPartResponse.builder().eTag(partETags.get(partNumber - 1)).build(); + }); + + // Mock complete response with final ETag + when(client.completeMultipartUpload(completeRequestCaptor.capture())).thenReturn( + CompleteMultipartUploadResponse.builder().eTag(finalETag).build() + ); + + // Create listener + @SuppressWarnings("unchecked") + ActionListener etagListener = mock(ActionListener.class); + + // Execute the multipart upload + final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); + blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, metadata, inputETag, etagListener); + + // Verify the create multipart request parameters including all configurations + final CreateMultipartUploadRequest createRequest = createRequestCaptor.getValue(); + assertEquals(bucketName, createRequest.bucket()); + assertEquals(blobPath.buildAsString() + blobName, createRequest.key()); + assertEquals(storageClass, createRequest.storageClass()); + assertEquals(cannedAccessControlList, createRequest.acl()); + assertEquals(metadata, createRequest.metadata()); + + // Verify server-side encryption is set + assertEquals(ServerSideEncryption.AES256, createRequest.serverSideEncryption()); + + // Verify upload part requests + List partRequests = uploadPartRequestCaptor.getAllValues(); + assertEquals(2, partRequests.size()); // We expect exactly 2 parts + + // Verify both parts + for (int i = 0; i < 2; i++) { + UploadPartRequest partRequest = partRequests.get(i); + assertEquals(bucketName, partRequest.bucket()); + assertEquals(blobPath.buildAsString() + blobName, partRequest.key()); + assertEquals(uploadId, partRequest.uploadId()); + assertEquals(Integer.valueOf(i + 1), partRequest.partNumber()); + assertEquals(partSize, partRequest.contentLength().longValue()); + + if (i == 0) { + RequestBody body = requestBodyCaptor.getAllValues().get(i); + try (InputStream is = body.contentStreamProvider().newStream()) { + assertNotNull(is); + assertTrue("Content stream should be available", is.available() > 0); + } + } + } + + // Verify complete request + CompleteMultipartUploadRequest completeRequest = completeRequestCaptor.getValue(); + assertEquals(bucketName, completeRequest.bucket()); + assertEquals(blobPath.buildAsString() + blobName, completeRequest.key()); + assertEquals(uploadId, completeRequest.uploadId()); + + // Verify ifMatch is on complete request - NOT create request + assertEquals(inputETag, completeRequest.ifMatch()); + + // Verify listener was called with exact final ETag + verify(etagListener).onResponse(finalETag); + verify(etagListener, never()).onFailure(any()); + + // Verify resources were properly closed + verify(clientReference).close(); + + } + + /** + * Tests that actual content is properly transmitted during conditional multipart upload + */ + public void testExecuteMultipartUploadIfEtagMatchesContentIntegrity() throws IOException { + // Setup test parameters + final String bucketName = randomAlphaOfLengthBetween(1, 10); + final String blobName = randomAlphaOfLengthBetween(1, 10); + final String inputETag = randomAlphaOfLengthBetween(8, 32); // Input ETag for condition + final String finalETag = randomAlphaOfLengthBetween(8, 32); // Final ETag returned + final String uploadId = randomAlphaOfLengthBetween(10, 20); + + final BlobPath blobPath = new BlobPath(); + + // Setup for multipart upload with known content + final int partCount = 3; + final long partSize = ByteSizeUnit.MB.toBytes(5); // 5MB + final long blobSize = partSize * partCount; + + // Create verifiable content - use a pattern that's different for each part + final byte[] blobContent = new byte[(int) blobSize]; + Random random = new Random(0); // Fixed seed for reproducibility + random.nextBytes(blobContent); + + // Mock S3BlobStore + final S3BlobStore blobStore = mock(S3BlobStore.class); + + // FIX 1: Replace mock StatsMetricPublisher with a real instance + final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); + + when(blobStore.bucket()).thenReturn(bucketName); + when(blobStore.bufferSizeInBytes()).thenReturn(partSize); + when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); + when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); + when(blobStore.serverSideEncryption()).thenReturn(false); + when(blobStore.isUploadRetryEnabled()).thenReturn(false); + + final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); + + // Mock S3 client + final S3Client client = mock(S3Client.class); + final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); + when(blobStore.clientReference()).thenReturn(clientReference); + when(clientReference.get()).thenReturn(client); + + // Setup argument captors with focus on request body + final ArgumentCaptor createRequestCaptor = ArgumentCaptor.forClass( + CreateMultipartUploadRequest.class + ); + final ArgumentCaptor completeRequestCaptor = ArgumentCaptor.forClass( + CompleteMultipartUploadRequest.class + ); + final ArgumentCaptor requestBodyCaptor = ArgumentCaptor.forClass(RequestBody.class); + + // Mock responses + when(client.createMultipartUpload(createRequestCaptor.capture())).thenReturn( + CreateMultipartUploadResponse.builder().uploadId(uploadId).build() + ); + + // Content verification will happen in this mock + final List capturedPartContents = new ArrayList<>(); + + when(client.uploadPart(any(UploadPartRequest.class), requestBodyCaptor.capture())).thenAnswer(invocation -> { + // Capture content from request body for later verification + RequestBody requestBody = requestBodyCaptor.getValue(); + try (InputStream contentStream = requestBody.contentStreamProvider().newStream()) { + byte[] partContent = contentStream.readAllBytes(); + capturedPartContents.add(partContent); + } + return UploadPartResponse.builder().eTag("etag-for-part").build(); + }); + + // Add final ETag to complete response + when(client.completeMultipartUpload(completeRequestCaptor.capture())).thenReturn( + CompleteMultipartUploadResponse.builder().eTag(finalETag).build() + ); + + // Create listener + @SuppressWarnings("unchecked") + ActionListener etagListener = mock(ActionListener.class); + + // FIX 2: Handle exceptions properly - in this case we expect successful execution + // If we were expecting an exception, we would use expectThrows instead + final ByteArrayInputStream inputStream = new ByteArrayInputStream(blobContent); + blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, null, inputETag, etagListener); + + // Verify key request parameters + final CreateMultipartUploadRequest createRequest = createRequestCaptor.getValue(); + assertEquals(bucketName, createRequest.bucket()); + assertEquals(blobPath.buildAsString() + blobName, createRequest.key()); + + final CompleteMultipartUploadRequest completeRequest = completeRequestCaptor.getValue(); + assertEquals(inputETag, completeRequest.ifMatch()); + assertEquals(bucketName, completeRequest.bucket()); + assertEquals(blobPath.buildAsString() + blobName, completeRequest.key()); + assertEquals(uploadId, completeRequest.uploadId()); + + assertEquals(partCount, capturedPartContents.size()); + + byte[] reassembledContent = new byte[(int) blobSize]; + int offset = 0; + for (byte[] partContent : capturedPartContents) { + System.arraycopy(partContent, 0, reassembledContent, offset, partContent.length); + offset += partContent.length; + } + + assertArrayEquals("Uploaded content should match original content", blobContent, reassembledContent); + + verify(etagListener).onResponse(finalETag); + verify(etagListener, never()).onFailure(any()); + + verify(clientReference).close(); + } + + /** + * Tests that multipart uploads enforce minimum and maximum size boundaries + */ + public void testExecuteMultipartUploadIfEtagMatchesSizeValidation() { + final S3BlobStore blobStore = mock(S3BlobStore.class); + final S3BlobContainer blobContainer = new S3BlobContainer(mock(BlobPath.class), blobStore); + final String blobName = randomAlphaOfLengthBetween(1, 10); + final String inputETag = randomAlphaOfLengthBetween(8, 32); + final String finalETag = randomAlphaOfLengthBetween(8, 32); + + // Use mocked listeners to verify interactions + @SuppressWarnings("unchecked") + ActionListener invalidSizeListener = mock(ActionListener.class); + + // Test case 1: Below minimum size (4.9MB) + { + final long tooSmallSize = ByteSizeUnit.MB.toBytes(5) - 1024; // Just below 5MB + + final IllegalArgumentException tooSmallException = expectThrows( + IllegalArgumentException.class, + () -> blobContainer.executeMultipartUploadIfEtagMatches( + blobStore, + blobName, + new ByteArrayInputStream(new byte[0]), + tooSmallSize, + null, + inputETag, + invalidSizeListener + ) + ); + + assertTrue(tooSmallException.getMessage().contains("can't be smaller than")); + // Listener should never be called when validation fails immediately + verify(invalidSizeListener, never()).onResponse(any()); + verify(invalidSizeListener, never()).onFailure(any()); + } + + // Test case 2: Above maximum size (5TB + 1 byte) + { + final long tooLargeSize = ByteSizeUnit.TB.toBytes(5) + 1; + + final IllegalArgumentException tooLargeException = expectThrows( + IllegalArgumentException.class, + () -> blobContainer.executeMultipartUploadIfEtagMatches( + blobStore, + blobName, + new ByteArrayInputStream(new byte[0]), + tooLargeSize, + null, + inputETag, + invalidSizeListener + ) + ); + + assertTrue(tooLargeException.getMessage().contains("can't be larger than")); + verify(invalidSizeListener, never()).onResponse(any()); + verify(invalidSizeListener, never()).onFailure(any()); + } + + // Setup mocks for valid size tests + final S3Client client = mock(S3Client.class); + final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); + final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); + + when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); + when(blobStore.clientReference()).thenReturn(clientReference); + when(clientReference.get()).thenReturn(client); + when(blobStore.bucket()).thenReturn("test-bucket"); + + // Configure buffer size to allow valid tests + when(blobStore.bufferSizeInBytes()).thenReturn(ByteSizeUnit.MB.toBytes(5)); + + // Add argument captors for request validation + ArgumentCaptor completeRequestCaptor = ArgumentCaptor.forClass( + CompleteMultipartUploadRequest.class + ); + + // Setup success responses + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( + CreateMultipartUploadResponse.builder().uploadId("test-upload-id").build() + ); + + when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenReturn( + UploadPartResponse.builder().eTag("test-etag").build() + ); + + when(client.completeMultipartUpload(completeRequestCaptor.capture())).thenReturn( + CompleteMultipartUploadResponse.builder().eTag(finalETag).build() + ); + + // Test case 3: Exactly minimum size (5MB) + { + final long exactMinimumSize = ByteSizeUnit.MB.toBytes(5); + @SuppressWarnings("unchecked") + ActionListener validSizeListener = mock(ActionListener.class); + + // Use a stream that returns zeros but doesn't allocate full memory + InputStream zeroStream = new InputStream() { + long remaining = exactMinimumSize; + + @Override + public int read() { + if (remaining > 0) { + remaining--; + return 0; + } else { + return -1; + } + } + + @Override + public int read(byte[] b, int off, int len) { + if (remaining <= 0) { + return -1; + } + int toRead = (int) Math.min(len, remaining); + Arrays.fill(b, off, off + toRead, (byte) 0); + remaining -= toRead; + return toRead; + } + }; + + try { + // This should not throw an exception + blobContainer.executeMultipartUploadIfEtagMatches( + blobStore, + blobName, + zeroStream, + exactMinimumSize, + null, + inputETag, + validSizeListener + ); + + // Verify that listener was called with success + verify(validSizeListener).onResponse(finalETag); + verify(validSizeListener, never()).onFailure(any()); + + // Verify resources were closed + verify(clientReference).close(); + + // Verify if-match condition was applied properly + CompleteMultipartUploadRequest completeRequest = completeRequestCaptor.getValue(); + assertEquals(inputETag, completeRequest.ifMatch()); + + } catch (IOException e) { + fail("Should not throw exception for exact minimum size: " + e); + } + } + + // Reset for next test case + reset(clientReference); + when(blobStore.clientReference()).thenReturn(clientReference); + when(clientReference.get()).thenReturn(client); + + // Test case 4: Valid larger size + { + // Use a moderate size for test purposes (10MB) + final long testSize = ByteSizeUnit.MB.toBytes(10); + @SuppressWarnings("unchecked") + ActionListener validSizeListener = mock(ActionListener.class); + + // Use a stream that returns zeros without allocating full memory + InputStream zeroStream = new InputStream() { + long remaining = testSize; + + @Override + public int read() { + if (remaining > 0) { + remaining--; + return 0; + } else { + return -1; + } + } + + @Override + public int read(byte[] b, int off, int len) { + if (remaining <= 0) { + return -1; + } + int toRead = (int) Math.min(len, remaining); + Arrays.fill(b, off, off + toRead, (byte) 0); + remaining -= toRead; + return toRead; + } + }; + + try { + blobContainer.executeMultipartUploadIfEtagMatches( + blobStore, + blobName, + zeroStream, + testSize, + null, + inputETag, + validSizeListener + ); + + // Verify that listener was called with success + verify(validSizeListener).onResponse(finalETag); + verify(validSizeListener, never()).onFailure(any()); + + // Verify resources were closed + verify(clientReference).close(); + + // Verify if-match condition was applied properly + CompleteMultipartUploadRequest completeRequest = completeRequestCaptor.getValue(); + assertEquals(inputETag, completeRequest.ifMatch()); + + } catch (IOException e) { + fail("Should not fail with size validation error: " + e); + } + } + } + + /** + * Tests handling of 412 Precondition Failed errors when ETag doesn't match + */ + public void testExecuteMultipartUploadIfEtagMatchesPreconditionFailed() { + // Setup test parameters + final String bucketName = randomAlphaOfLengthBetween(1, 10); + final String blobName = randomAlphaOfLengthBetween(1, 10); + final String eTag = randomAlphaOfLengthBetween(8, 32); + + final BlobPath blobPath = new BlobPath(); + final long partSize = ByteSizeUnit.MB.toBytes(5); // 5MB + final long blobSize = partSize * 2; // Two parts for simplicity + + // Mock S3BlobStore + final S3BlobStore blobStore = mock(S3BlobStore.class); + final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); + when(blobStore.bucket()).thenReturn(bucketName); + when(blobStore.bufferSizeInBytes()).thenReturn(partSize); + when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); + when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); + when(blobStore.serverSideEncryption()).thenReturn(false); + when(blobStore.isUploadRetryEnabled()).thenReturn(false); + + final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); + + // Mock S3 client and references + final S3Client client = mock(S3Client.class); + final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); + final AmazonS3Reference abortClientReference = mock(AmazonS3Reference.class); + + // Setup client reference behavior + when(blobStore.clientReference()).thenReturn(clientReference).thenReturn(abortClientReference); + when(clientReference.get()).thenReturn(client); + when(abortClientReference.get()).thenReturn(client); + + // Create 412 Precondition Failed exception + S3Exception preconditionFailedException = (S3Exception) S3Exception.builder() + .message("Precondition Failed") + .statusCode(412) + .build(); + + // Setup success for initial upload stages + final String uploadId = randomAlphaOfLengthBetween(10, 20); + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( + CreateMultipartUploadResponse.builder().uploadId(uploadId).build() + ); + + when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenReturn( + UploadPartResponse.builder().eTag("part-etag").build() + ); + + // Fail on complete with 412 + when(client.completeMultipartUpload(any(CompleteMultipartUploadRequest.class))).thenThrow(preconditionFailedException); + + // Setup abort response + when(client.abortMultipartUpload(any(AbortMultipartUploadRequest.class))).thenReturn( + AbortMultipartUploadResponse.builder().build() + ); + + // Execute the multipart upload + final AtomicReference capturedException = new AtomicReference<>(); + ActionListener etagListener = ActionListener.wrap( + r -> fail("Should have failed with precondition failure"), + capturedException::set + ); + + final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); + + IOException ioException = expectThrows( + IOException.class, + () -> blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, null, eTag, etagListener) + ); + + // Verify exception handling + assertEquals("Unable to upload object [" + blobName + "] due to ETag mismatch", ioException.getMessage()); + assertEquals(preconditionFailedException, ioException.getCause()); + + Exception exception = capturedException.get(); + assertNotNull("Expected an exception to be captured", exception); + assertTrue("Exception should be an OpenSearchException", exception instanceof OpenSearchException); + assertEquals("stale_primary_shard", ((OpenSearchException) exception).getMessage()); + + // Verify workflow + verify(client).createMultipartUpload(any(CreateMultipartUploadRequest.class)); + verify(client).completeMultipartUpload(any(CompleteMultipartUploadRequest.class)); + verify(client).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); + + // Verify both client references were closed + verify(clientReference).close(); + verify(abortClientReference).close(); + } + + /** + * Tests handling of various exception types during multipart uploads + */ + public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { + // Setup basic test parameters + final String bucketName = randomAlphaOfLengthBetween(1, 10); + final String blobName = randomAlphaOfLengthBetween(1, 10); + final String eTag = randomAlphaOfLengthBetween(8, 32); + + final BlobPath blobPath = new BlobPath(); + final long partSize = ByteSizeUnit.MB.toBytes(5); + final long blobSize = partSize * 2; + + // Create test metadata + Map metadata = new HashMap<>(); + metadata.put("key1", "value1"); + + // Define test cases with minimal necessary parameters + // [Exception type, error phase (0=create, 1=upload, 2=complete), status code] + Object[][] testCases = { + // Key test cases for each error path + { "S3Exception", 0, 403 }, // Access denied during create + { "S3Exception", 1, 404 }, // Not found during upload + { "S3Exception", 2, 412 }, // Conditional check failure during complete + { "SdkException", 1, 0 }, // SDK error during upload + }; + + for (Object[] testCase : testCases) { + String exceptionType = (String) testCase[0]; + int errorPhase = (int) testCase[1]; + int statusCode = (int) testCase[2]; + + // Create fresh mocks for each test case + final S3BlobStore blobStore = mock(S3BlobStore.class); + final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); + + // Basic configuration + when(blobStore.bucket()).thenReturn(bucketName); + when(blobStore.bufferSizeInBytes()).thenReturn(partSize); + when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); + when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); + when(blobStore.serverSideEncryption()).thenReturn(false); + + final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); + + // Create client references + final S3Client client = mock(S3Client.class); + final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); + final AmazonS3Reference abortClientReference = mock(AmazonS3Reference.class); + + // Setup client reference behavior + when(blobStore.clientReference()).thenReturn(clientReference).thenReturn(abortClientReference); + when(clientReference.get()).thenReturn(client); + when(abortClientReference.get()).thenReturn(client); + + // Create the appropriate exception + Exception testException; + if ("S3Exception".equals(exceptionType)) { + testException = (S3Exception) S3Exception.builder() + .message("S3 Error with status code " + statusCode) + .statusCode(statusCode) + .build(); + } else { + testException = SdkException.builder().message("SDK Error occurred").build(); + } + + // Setup responses for each phase + final String uploadId = "test-upload-id"; + if (errorPhase >= 1) { + // Success for create if error happens later + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( + CreateMultipartUploadResponse.builder().uploadId(uploadId).build() + ); + } else { + // Error during create + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenThrow(testException); + } + + if (errorPhase >= 2) { + // Success for upload parts if error happens at complete + when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenReturn( + UploadPartResponse.builder().eTag("part-etag").build() + ); + } else if (errorPhase == 1) { + // Error during upload + when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenThrow(testException); + } + + if (errorPhase == 2) { + // Error during complete + when(client.completeMultipartUpload(any(CompleteMultipartUploadRequest.class))).thenThrow(testException); + } + + // Setup abort response + when(client.abortMultipartUpload(any(AbortMultipartUploadRequest.class))).thenReturn( + AbortMultipartUploadResponse.builder().build() + ); + + // Capture listener exception + final AtomicReference capturedException = new AtomicReference<>(); + ActionListener etagListener = ActionListener.wrap( + r -> fail("Should have failed with exception"), + capturedException::set + ); + + // Execute with expectThrows + final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); + IOException exception = expectThrows( + IOException.class, + () -> blobContainer.executeMultipartUploadIfEtagMatches( + blobStore, + blobName, + inputStream, + blobSize, + metadata, + eTag, + etagListener + ) + ); + + // Verify thrown exception has the correct cause + assertEquals(testException, exception.getCause()); + + // Verify listener behavior + Exception listenerException = capturedException.get(); + assertNotNull("Expected a listener exception", listenerException); + + // Special handling for 412 case + if ("S3Exception".equals(exceptionType) && statusCode == 412) { + assertTrue(listenerException instanceof OpenSearchException); + assertEquals("stale_primary_shard", ((OpenSearchException) listenerException).getMessage()); + } else { + assertTrue(listenerException instanceof IOException); + } + + // Verify resource cleanup + verify(clientReference).close(); + + // Verify abort called when appropriate + if (errorPhase >= 1) { + verify(client).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); + verify(abortClientReference).close(); + } else { + verify(client, never()).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); + } + } + } + + /** + * Tests handling of generic SdkException scenarios during multipart uploads + */ + public void testExecuteMultipartUploadIfEtagMatchesSdkException() { + // Test parameters + final String bucketName = randomAlphaOfLengthBetween(1, 10); + final String blobName = randomAlphaOfLengthBetween(1, 10); + final String eTag = randomAlphaOfLengthBetween(8, 32); + + final BlobPath blobPath = new BlobPath(); + final long partSize = ByteSizeUnit.MB.toBytes(5); + final long blobSize = partSize * 2; + + // Define test cases: [scenario name, error stage, should abort fail] + Object[][] testScenarios = { + { "initialization error", 0, false }, // SdkException during createMultipartUpload + { "part upload error", 1, false }, // SdkException during uploadPart with successful abort + { "abort failure", 1, true } // SdkException during uploadPart with abort failure + }; + + for (Object[] scenario : testScenarios) { + String scenarioName = (String) scenario[0]; + int errorStage = (int) scenario[1]; + boolean abortFails = (boolean) scenario[2]; + + // Setup fresh mocks for each scenario + final S3BlobStore blobStore = mock(S3BlobStore.class); + final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); + when(blobStore.bucket()).thenReturn(bucketName); + when(blobStore.bufferSizeInBytes()).thenReturn(partSize); + when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); + when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); + + final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); + + // Mock S3 client + final S3Client client = mock(S3Client.class); + final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); + final AmazonS3Reference abortClientReference = mock(AmazonS3Reference.class); + when(blobStore.clientReference()).thenReturn(clientReference).thenReturn(abortClientReference); + when(clientReference.get()).thenReturn(client); + when(abortClientReference.get()).thenReturn(client); + + // Create primary exception + SdkException primaryException = SdkException.builder().message("SDK error during " + scenarioName).build(); + + // Setup responses based on scenario + final String uploadId = "test-upload-id"; + if (errorStage == 0) { + // Error during initialization + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenThrow(primaryException); + } else { + // Success for initialization + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( + CreateMultipartUploadResponse.builder().uploadId(uploadId).build() + ); + + // Error during part upload + when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenThrow(primaryException); + } + + // Configure abort behavior with argument captor + ArgumentCaptor abortCaptor = ArgumentCaptor.forClass(AbortMultipartUploadRequest.class); + + if (abortFails) { + SdkException abortException = SdkException.builder().message("Abort failure").build(); + when(client.abortMultipartUpload(abortCaptor.capture())).thenThrow(abortException); + } else { + when(client.abortMultipartUpload(abortCaptor.capture())).thenReturn(AbortMultipartUploadResponse.builder().build()); + } + + // Capture listener exception + final AtomicReference capturedException = new AtomicReference<>(); + ActionListener etagListener = ActionListener.wrap( + r -> fail("Should have failed with SdkException"), + capturedException::set + ); + + // Execute with expectThrows + final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); + IOException exception = expectThrows( + IOException.class, + () -> blobContainer.executeMultipartUploadIfEtagMatches( + blobStore, + blobName, + inputStream, + blobSize, + null, + eTag, + etagListener + ) + ); + + // Verify thrown exception + assertEquals("Original SdkException should be preserved as cause", primaryException, exception.getCause()); + + // Verify listener exception + Exception listenerException = capturedException.get(); + assertNotNull("Expected an exception", listenerException); + assertTrue("Exception should be IOException", listenerException instanceof IOException); + + // Verify main client reference was closed + verify(clientReference).close(); + + // Verify abort behavior + if (errorStage > 0) { + // Should attempt abort if upload ID exists + verify(client).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); + verify(abortClientReference).close(); + + // Verify abort request parameters + if (!abortCaptor.getAllValues().isEmpty()) { + AbortMultipartUploadRequest abortRequest = abortCaptor.getValue(); + assertEquals("Abort request should have correct upload ID", uploadId, abortRequest.uploadId()); + assertEquals("Abort request should have correct bucket", bucketName, abortRequest.bucket()); + assertEquals("Abort request should have correct key", blobPath.buildAsString() + blobName, abortRequest.key()); + } + } else { + // No abort if no upload ID + verify(client, never()).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); + } + } + } + + /** + * Tests that client references are properly managed (acquired and closed) in various scenarios + */ + public void testExecuteMultipartUploadIfEtagMatchesResourceManagement() throws IOException { + // Setup test parameters + final String bucketName = randomAlphaOfLengthBetween(1, 10); + final String blobName = randomAlphaOfLengthBetween(1, 10); + final String eTag = randomAlphaOfLengthBetween(8, 32); + final String uploadId = randomAlphaOfLengthBetween(10, 20); + + final BlobPath blobPath = new BlobPath(); + final long partSize = ByteSizeUnit.MB.toBytes(5); + final long blobSize = partSize * 2; + + // Define scenarios to test different resource management cases + enum ResourceScenario { + SUCCESS_PATH, // Complete successful upload + INIT_FAILURE, // Failure during createMultipartUpload + PART_UPLOAD_FAILURE, // Failure during uploadPart + COMPLETION_FAILURE, // Failure during completeMultipartUpload + ABORT_FAILURE // Failure during abortMultipartUpload + } + + for (ResourceScenario scenario : ResourceScenario.values()) { + // Create fresh mocks for each scenario + final S3BlobStore blobStore = mock(S3BlobStore.class); + final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); + when(blobStore.bucket()).thenReturn(bucketName); + when(blobStore.bufferSizeInBytes()).thenReturn(partSize); + when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); + when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); + + // S3 client and references + final S3Client client = mock(S3Client.class); + + // The implementation always creates separate references for the main operation and abort + AmazonS3Reference primaryClientReference = mock(AmazonS3Reference.class); + AmazonS3Reference abortClientReference = mock(AmazonS3Reference.class); + + when(primaryClientReference.get()).thenReturn(client); + when(abortClientReference.get()).thenReturn(client); + + // Configure clientReference behavior based on implementation pattern + when(blobStore.clientReference()).thenReturn(primaryClientReference).thenReturn(abortClientReference); + + final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); + + // Create exceptions with distinctive messages for each scenario + SdkException stageException = SdkException.builder().message("Failure during " + scenario.name()).build(); + + // Setup argument captor for abort request validation + ArgumentCaptor abortCaptor = ArgumentCaptor.forClass(AbortMultipartUploadRequest.class); + + // Configure upload behavior based on scenario + switch (scenario) { + case SUCCESS_PATH: + // Successful path should close reference after completion + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( + CreateMultipartUploadResponse.builder().uploadId(uploadId).build() + ); + when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenReturn( + UploadPartResponse.builder().eTag("test-etag").build() + ); + when(client.completeMultipartUpload(any(CompleteMultipartUploadRequest.class))).thenReturn( + CompleteMultipartUploadResponse.builder().eTag("final-etag").build() + ); + break; + + case INIT_FAILURE: + // Initialization failure should close reference immediately + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenThrow(stageException); + break; + + case PART_UPLOAD_FAILURE: + // Part upload failure should close primary reference and abort reference + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( + CreateMultipartUploadResponse.builder().uploadId(uploadId).build() + ); + when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenThrow(stageException); + when(client.abortMultipartUpload(abortCaptor.capture())).thenReturn(AbortMultipartUploadResponse.builder().build()); + break; + + case COMPLETION_FAILURE: + // Completion failure should close primary reference and abort reference + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( + CreateMultipartUploadResponse.builder().uploadId(uploadId).build() + ); + when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenReturn( + UploadPartResponse.builder().eTag("test-etag").build() + ); + when(client.completeMultipartUpload(any(CompleteMultipartUploadRequest.class))).thenThrow(stageException); + when(client.abortMultipartUpload(abortCaptor.capture())).thenReturn(AbortMultipartUploadResponse.builder().build()); + break; + + case ABORT_FAILURE: + // Even if abort fails, all references should be closed + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( + CreateMultipartUploadResponse.builder().uploadId(uploadId).build() + ); + when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenThrow(stageException); + when(client.abortMultipartUpload(abortCaptor.capture())).thenThrow( + SdkException.builder().message("Abort failure").build() + ); + break; + } + + // Create action listener to capture callbacks + @SuppressWarnings("unchecked") + ActionListener etagListener = mock(ActionListener.class); + + final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); + + // Execute the appropriate test based on scenario + if (scenario == ResourceScenario.SUCCESS_PATH) { + // Success path should not throw exceptions + blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, null, eTag, etagListener); + + // Verify listener success callback + verify(etagListener).onResponse("final-etag"); + verify(etagListener, never()).onFailure(any()); + + // Success path uses one client reference exactly once + verify(blobStore, times(1)).clientReference(); + verify(primaryClientReference).close(); + } else { + // For all failure scenarios, expect exception + IOException exception = expectThrows( + IOException.class, + () -> blobContainer.executeMultipartUploadIfEtagMatches( + blobStore, + blobName, + inputStream, + blobSize, + null, + eTag, + etagListener + ) + ); + + // Verify exception has the correct cause + assertEquals("Exception cause should be the original exception", stageException, exception.getCause()); + + // Verify listener failure callback + verify(etagListener).onFailure(any(Exception.class)); + verify(etagListener, never()).onResponse(any()); + + // Primary reference should always be closed + verify(primaryClientReference).close(); + + if (scenario != ResourceScenario.INIT_FAILURE) { + // For non-init failures, verify abort is called + verify(blobStore, times(2)).clientReference(); // Main + abort + verify(client).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); + verify(abortClientReference).close(); + + // Verify abort request parameters + if (!abortCaptor.getAllValues().isEmpty()) { + AbortMultipartUploadRequest abortRequest = abortCaptor.getValue(); + assertEquals("Upload ID should match", uploadId, abortRequest.uploadId()); + assertEquals("Bucket should match", bucketName, abortRequest.bucket()); + assertEquals("Key should match", blobPath.buildAsString() + blobName, abortRequest.key()); + } + } else { + // Init failure should not use abort + verify(blobStore, times(1)).clientReference(); // Main only + verify(client, never()).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); + } + } + } + } + public void testInitCannedACL() { String[] aclList = new String[] { "private", From 2c31c83de99b630c9954a837fef6aa9da97cbd7a Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Mon, 28 Apr 2025 05:16:27 +0530 Subject: [PATCH 02/20] git commit -m "Implementation for : executeMultipartUploadIfEtagMatches" --- .../repositories/s3/S3BlobContainer.java | 167 ++++++++++++++++++ 1 file changed, 167 insertions(+) diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java index e83ca97b385f0..bb3d4eafaad12 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java @@ -40,6 +40,7 @@ import software.amazon.awssdk.services.s3.model.AbortMultipartUploadRequest; import software.amazon.awssdk.services.s3.model.CommonPrefix; import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest; +import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse; import software.amazon.awssdk.services.s3.model.CompletedMultipartUpload; import software.amazon.awssdk.services.s3.model.CompletedPart; import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest; @@ -53,6 +54,7 @@ import software.amazon.awssdk.services.s3.model.NoSuchKeyException; import software.amazon.awssdk.services.s3.model.ObjectAttributes; import software.amazon.awssdk.services.s3.model.PutObjectRequest; +import software.amazon.awssdk.services.s3.model.S3Exception; import software.amazon.awssdk.services.s3.model.ServerSideEncryption; import software.amazon.awssdk.services.s3.model.UploadPartRequest; import software.amazon.awssdk.services.s3.model.UploadPartResponse; @@ -63,6 +65,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.logging.log4j.message.ParameterizedMessage; +import org.opensearch.OpenSearchException; import org.opensearch.action.support.PlainActionFuture; import org.opensearch.common.Nullable; import org.opensearch.common.SetOnce; @@ -97,6 +100,7 @@ import java.io.InputStream; import java.util.ArrayList; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; @@ -509,6 +513,169 @@ private String buildKey(String blobName) { return keyPath + blobName; } + public void executeMultipartUploadIfEtagMatches( + final S3BlobStore blobStore, + final String blobName, + final InputStream input, + final long blobSize, + final Map metadata, + final String eTag, + final ActionListener etagListener + ) throws IOException { + + ensureMultiPartUploadSize(blobSize); + + final long partSize = blobStore.bufferSizeInBytes(); + final Tuple multiparts = numberOfMultiparts(blobSize, partSize); + if (multiparts.v1() > Integer.MAX_VALUE) { + throw new IllegalArgumentException("Too many multipart upload parts; consider a larger buffer size."); + } + final int nbParts = multiparts.v1().intValue(); + final long lastPartSize = multiparts.v2(); + assert blobSize == (((nbParts - 1) * partSize) + lastPartSize) : "blobSize does not match multipart sizes"; + // test + CreateMultipartUploadRequest.Builder createRequestBuilder = CreateMultipartUploadRequest.builder() + .bucket(blobStore.bucket()) + .key(blobName) + .storageClass(blobStore.getStorageClass()) + .acl(blobStore.getCannedACL()) + .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector)); + + if (metadata != null && !metadata.isEmpty()) { + createRequestBuilder.metadata(metadata); + } + if (blobStore.serverSideEncryption()) { + createRequestBuilder.serverSideEncryption(ServerSideEncryption.AES256); + } + + final CreateMultipartUploadRequest createMultipartUploadRequest = createRequestBuilder.build(); + final SetOnce uploadId = new SetOnce<>(); + final String bucketName = blobStore.bucket(); + boolean success = false; + + final InputStream requestInputStream = blobStore.isUploadRetryEnabled() + ? new BufferedInputStream(input, (int) (partSize + 1)) + : input; + + try (AmazonS3Reference clientReference = blobStore.clientReference()) { + uploadId.set( + SocketAccess.doPrivileged(() -> clientReference.get().createMultipartUpload(createMultipartUploadRequest).uploadId()) + ); + if (Strings.isEmpty(uploadId.get())) { + IOException exception = new IOException("Failed to initialize multipart upload for " + blobName); + etagListener.onFailure(exception); + throw exception; + } + + final List parts = new ArrayList<>(nbParts); + long bytesCount = 0; + + for (int i = 1; i <= nbParts; i++) { + long currentPartSize = (i < nbParts) ? partSize : lastPartSize; + final UploadPartRequest uploadPartRequest = UploadPartRequest.builder() + .bucket(bucketName) + .key(blobName) + .uploadId(uploadId.get()) + .partNumber(i) + .contentLength(currentPartSize) + .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector)) + .build(); + + bytesCount += currentPartSize; + + final UploadPartResponse uploadResponse = SocketAccess.doPrivileged( + () -> clientReference.get() + .uploadPart(uploadPartRequest, RequestBody.fromInputStream(requestInputStream, currentPartSize)) + ); + + String partETag = uploadResponse.eTag(); + if (partETag == null) { + IOException exception = new IOException( + String.format(Locale.ROOT, "S3 part upload for [%s] part [%d] returned null ETag", blobName, i) + ); + etagListener.onFailure(exception); + throw exception; + } + + parts.add(CompletedPart.builder().partNumber(i).eTag(partETag).build()); + } + + if (bytesCount != blobSize) { + IOException exception = new IOException( + String.format(Locale.ROOT, "Multipart upload for [%s] sent %d bytes; expected %d bytes", blobName, bytesCount, blobSize) + ); + etagListener.onFailure(exception); + throw exception; + } + + CompleteMultipartUploadRequest completeRequest = CompleteMultipartUploadRequest.builder() + .bucket(bucketName) + .key(blobName) + .uploadId(uploadId.get()) + .multipartUpload(CompletedMultipartUpload.builder().parts(parts).build()) + .ifMatch(eTag) + .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector)) + .build(); + + CompleteMultipartUploadResponse completeResponse = SocketAccess.doPrivileged( + () -> clientReference.get().completeMultipartUpload(completeRequest) + ); + + if (completeResponse.eTag() != null) { + success = true; + etagListener.onResponse(completeResponse.eTag()); + } else { + IOException exception = new IOException( + "S3 multipart upload for [" + blobName + "] returned null ETag, violating data integrity expectations" + ); + etagListener.onFailure(exception); + throw exception; + } + + } catch (S3Exception e) { + if (e.statusCode() == 412) { + etagListener.onFailure(new OpenSearchException("stale_primary_shard", e, "Precondition Failed : Etag Mismatch", blobName)); + throw new IOException("Unable to upload object [" + blobName + "] due to ETag mismatch", e); + } else { + IOException exception = new IOException( + String.format(Locale.ROOT, "S3 error during multipart upload [%s]: %s", blobName, e.getMessage()), + e + ); + etagListener.onFailure(exception); + throw exception; + } + } catch (SdkException e) { + IOException exception = new IOException(String.format(Locale.ROOT, "S3 multipart upload failed for [%s]", blobName), e); + etagListener.onFailure(exception); + throw exception; + } catch (Exception e) { + IOException exception = new IOException( + String.format(Locale.ROOT, "Unexpected error during multipart upload [%s]: %s", blobName, e.getMessage()), + e + ); + etagListener.onFailure(exception); + throw exception; + } finally { + if (!success && Strings.hasLength(uploadId.get())) { + AbortMultipartUploadRequest abortRequest = AbortMultipartUploadRequest.builder() + .bucket(bucketName) + .key(blobName) + .uploadId(uploadId.get()) + .build(); + try (AmazonS3Reference abortClient = blobStore.clientReference()) { + SocketAccess.doPrivilegedVoid(() -> abortClient.get().abortMultipartUpload(abortRequest)); + } catch (Exception abortException) { + logger.warn( + "Failed to abort incomplete multipart upload [{}] with ID [{}]. " + + "This may result in orphaned S3 data and charges.", + new Object[] { blobName, uploadId.get() }, + abortException + ); + } + } + } + } + /** * Uploads a blob using a single upload request */ From eb3778c24ef94aecb7606d64d6d47e5d88d51fc6 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Tue, 29 Apr 2025 03:12:09 +0530 Subject: [PATCH 03/20] Adds a constant HTTP_STATUS_PRECONDITION_FAILED with value 412 Signed-off-by: Tanishq Ranjan --- .../repositories/s3/S3BlobContainer.java | 3 +- .../s3/S3BlobStoreContainerTests.java | 250 ++---------------- 2 files changed, 31 insertions(+), 222 deletions(-) diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java index bb3d4eafaad12..55a170173c4ab 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java @@ -121,6 +121,7 @@ class S3BlobContainer extends AbstractBlobContainer implements AsyncMultiStreamB private final S3BlobStore blobStore; private final String keyPath; + public static final int HTTP_STATUS_PRECONDITION_FAILED = 412; S3BlobContainer(BlobPath path, S3BlobStore blobStore) { super(path); @@ -633,7 +634,7 @@ public void executeMultipartUploadIfEtagMatches( } } catch (S3Exception e) { - if (e.statusCode() == 412) { + if (e.statusCode() == HTTP_STATUS_PRECONDITION_FAILED) { etagListener.onFailure(new OpenSearchException("stale_primary_shard", e, "Precondition Failed : Etag Mismatch", blobName)); throw new IOException("Unable to upload object [" + blobName + "] due to ETag mismatch", e); } else { diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java index dba8044c19c80..4e9be32972e36 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java @@ -944,50 +944,39 @@ public void testNumberOfMultiparts() { assertNumberOfMultiparts(factor + 1, remaining, (size * factor) + remaining, size); } - /** - * Tests the basic success path for conditional multipart upload when ETag matches - */ public void testExecuteMultipartUploadIfEtagMatchesSuccess() throws IOException { - // Setup test parameters with randomized values final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); - final String inputETag = randomAlphaOfLengthBetween(8, 32); // Input ETag for conditional check - final String finalETag = randomAlphaOfLengthBetween(8, 32); // ETag returned by complete operation + final String inputETag = randomAlphaOfLengthBetween(8, 32); + final String finalETag = randomAlphaOfLengthBetween(8, 32); final String uploadId = randomAlphaOfLengthBetween(10, 20); - // Add metadata - ADDED final Map metadata = new HashMap<>(); metadata.put("key1", "value1"); metadata.put("key2", "value2"); - // Setup path final BlobPath blobPath = new BlobPath(); if (randomBoolean()) { IntStream.of(randomIntBetween(1, 5)).forEach(value -> blobPath.add("path_" + value)); } - // Calculate sizes for multipart upload - ensure minimum viable multipart size - final long partSize = ByteSizeUnit.MB.toBytes(5); // 5MB minimum part size + final long partSize = ByteSizeUnit.MB.toBytes(5); final int partCount = randomIntBetween(2, 5); final long lastPartSize = randomIntBetween(1, (int) partSize); final long blobSize = partSize * (partCount - 1) + lastPartSize; - // Mock S3BlobStore with randomized parameters - ADDED final S3BlobStore blobStore = mock(S3BlobStore.class); final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); when(blobStore.bucket()).thenReturn(bucketName); when(blobStore.bufferSizeInBytes()).thenReturn(partSize); when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); - // Randomize storage class - ADDED final StorageClass storageClass = randomFrom(StorageClass.values()); when(blobStore.getStorageClass()).thenReturn(storageClass); - // Randomize server-side encryption - ADDED final boolean serverSideEncryption = randomBoolean(); when(blobStore.serverSideEncryption()).thenReturn(serverSideEncryption); - // Randomize ACL - ADDED final ObjectCannedACL cannedAccessControlList = randomBoolean() ? randomFrom(ObjectCannedACL.values()) : null; if (cannedAccessControlList != null) { when(blobStore.getCannedACL()).thenReturn(cannedAccessControlList); @@ -998,13 +987,11 @@ public void testExecuteMultipartUploadIfEtagMatchesSuccess() throws IOException final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); - // Mock S3 client with argument captors to verify requests final S3Client client = mock(S3Client.class); final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); when(blobStore.clientReference()).thenReturn(clientReference); when(clientReference.get()).thenReturn(client); - // Setup argument captors for request validation final ArgumentCaptor createRequestCaptor = ArgumentCaptor.forClass( CreateMultipartUploadRequest.class ); @@ -1014,38 +1001,31 @@ public void testExecuteMultipartUploadIfEtagMatchesSuccess() throws IOException CompleteMultipartUploadRequest.class ); - // Mock createMultipartUpload response with uploadId when(client.createMultipartUpload(createRequestCaptor.capture())).thenReturn( CreateMultipartUploadResponse.builder().uploadId(uploadId).build() ); - // Generate unique ETags for each part final List partETags = new ArrayList<>(); for (int i = 0; i < partCount; i++) { partETags.add("etag-part-" + (i + 1)); } - // Mock uploadPart responses with the appropriate ETags when(client.uploadPart(uploadPartRequestCaptor.capture(), requestBodyCaptor.capture())).thenAnswer(invocation -> { UploadPartRequest request = (UploadPartRequest) invocation.getArguments()[0]; int partNumber = request.partNumber(); return UploadPartResponse.builder().eTag(partETags.get(partNumber - 1)).build(); }); - // Mock successful completion response WITH final ETag - FIXED when(client.completeMultipartUpload(completeRequestCaptor.capture())).thenReturn( CompleteMultipartUploadResponse.builder().eTag(finalETag).build() ); - // Create action listener to verify callback @SuppressWarnings("unchecked") ActionListener etagListener = mock(ActionListener.class); - // Execute the multipart upload final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, metadata, inputETag, etagListener); - // Verify the create multipart request parameters final CreateMultipartUploadRequest createRequest = createRequestCaptor.getValue(); assertEquals(bucketName, createRequest.bucket()); assertEquals(blobPath.buildAsString() + blobName, createRequest.key()); @@ -1053,88 +1033,69 @@ public void testExecuteMultipartUploadIfEtagMatchesSuccess() throws IOException assertEquals(cannedAccessControlList, createRequest.acl()); assertEquals(metadata, createRequest.metadata()); - // Verify server-side encryption if enabled - ADDED if (serverSideEncryption) { assertEquals(ServerSideEncryption.AES256, createRequest.serverSideEncryption()); } - // Verify part upload requests List partRequests = uploadPartRequestCaptor.getAllValues(); assertEquals(partCount, partRequests.size()); - // Verify request bodies for content (partial content verification) - ADDED List requestBodies = requestBodyCaptor.getAllValues(); assertEquals(partCount, requestBodies.size()); - // Check each part request for (int i = 0; i < partCount; i++) { UploadPartRequest partRequest = partRequests.get(i); assertEquals(bucketName, partRequest.bucket()); assertEquals(blobPath.buildAsString() + blobName, partRequest.key()); assertEquals(uploadId, partRequest.uploadId()); assertEquals(Integer.valueOf(i + 1), partRequest.partNumber()); - // Verify part size - last part can be different long expectedPartSize = (i < partCount - 1) ? partSize : lastPartSize; assertEquals(expectedPartSize, partRequest.contentLength().longValue()); - // Verify content for first part only (partial content verification) - MODIFIED if (i == 0) { RequestBody body = requestBodies.get(i); try (InputStream is = body.contentStreamProvider().newStream()) { - // Just verify the stream exists assertNotNull(is); } } } - // Verify complete request CompleteMultipartUploadRequest completeRequest = completeRequestCaptor.getValue(); assertEquals(bucketName, completeRequest.bucket()); assertEquals(blobPath.buildAsString() + blobName, completeRequest.key()); assertEquals(uploadId, completeRequest.uploadId()); - // Verify ifMatch on complete request (NOT create request) - FIXED assertEquals(inputETag, completeRequest.ifMatch()); - // Verify listener was called with the FINAL ETag (not last part's ETag) - FIXED verify(etagListener).onResponse(finalETag); verify(etagListener, never()).onFailure(any()); - // Verify resources were properly closed verify(clientReference).close(); } - /** - * Tests conditional multipart upload with various configuration combinations - */ public void testExecuteMultipartUploadIfEtagMatchesWithMetadataAndSSE() throws IOException { - // Setup test parameters final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String inputETag = randomAlphaOfLengthBetween(8, 32); - final String finalETag = randomAlphaOfLengthBetween(8, 32); // Added distinct final ETag + final String finalETag = randomAlphaOfLengthBetween(8, 32); final String uploadId = randomAlphaOfLengthBetween(10, 20); - // Create metadata map final Map metadata = new HashMap<>(); metadata.put("key1", "value1"); metadata.put("key2", "value2"); final BlobPath blobPath = new BlobPath(); - // Setup for a minimal viable multipart upload final long partSize = ByteSizeUnit.MB.toBytes(5); - final long blobSize = partSize * 2; // Two parts + final long blobSize = partSize * 2; - // Mock S3BlobStore with various configurations enabled final S3BlobStore blobStore = mock(S3BlobStore.class); final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); when(blobStore.bucket()).thenReturn(bucketName); when(blobStore.bufferSizeInBytes()).thenReturn(partSize); when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); - // Setup random configurations - final boolean serverSideEncryption = true; // Force enabled to test this case + final boolean serverSideEncryption = true; when(blobStore.serverSideEncryption()).thenReturn(serverSideEncryption); final StorageClass storageClass = randomFrom(StorageClass.values()); @@ -1148,13 +1109,11 @@ public void testExecuteMultipartUploadIfEtagMatchesWithMetadataAndSSE() throws I final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); - // Mock S3 client and capture requests final S3Client client = mock(S3Client.class); final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); when(blobStore.clientReference()).thenReturn(clientReference); when(clientReference.get()).thenReturn(client); - // Setup argument captors final ArgumentCaptor createRequestCaptor = ArgumentCaptor.forClass( CreateMultipartUploadRequest.class ); @@ -1164,35 +1123,28 @@ public void testExecuteMultipartUploadIfEtagMatchesWithMetadataAndSSE() throws I CompleteMultipartUploadRequest.class ); - // Mock responses when(client.createMultipartUpload(createRequestCaptor.capture())).thenReturn( CreateMultipartUploadResponse.builder().uploadId(uploadId).build() ); - // Generate part ETags final List partETags = List.of("etag-part-1", "etag-part-2"); - // Mock part upload responses when(client.uploadPart(uploadPartRequestCaptor.capture(), requestBodyCaptor.capture())).thenAnswer(invocation -> { UploadPartRequest request = (UploadPartRequest) invocation.getArguments()[0]; int partNumber = request.partNumber(); return UploadPartResponse.builder().eTag(partETags.get(partNumber - 1)).build(); }); - // Mock complete response with final ETag when(client.completeMultipartUpload(completeRequestCaptor.capture())).thenReturn( CompleteMultipartUploadResponse.builder().eTag(finalETag).build() ); - // Create listener @SuppressWarnings("unchecked") ActionListener etagListener = mock(ActionListener.class); - // Execute the multipart upload final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, metadata, inputETag, etagListener); - // Verify the create multipart request parameters including all configurations final CreateMultipartUploadRequest createRequest = createRequestCaptor.getValue(); assertEquals(bucketName, createRequest.bucket()); assertEquals(blobPath.buildAsString() + blobName, createRequest.key()); @@ -1200,14 +1152,11 @@ public void testExecuteMultipartUploadIfEtagMatchesWithMetadataAndSSE() throws I assertEquals(cannedAccessControlList, createRequest.acl()); assertEquals(metadata, createRequest.metadata()); - // Verify server-side encryption is set assertEquals(ServerSideEncryption.AES256, createRequest.serverSideEncryption()); - // Verify upload part requests List partRequests = uploadPartRequestCaptor.getAllValues(); - assertEquals(2, partRequests.size()); // We expect exactly 2 parts + assertEquals(2, partRequests.size()); - // Verify both parts for (int i = 0; i < 2; i++) { UploadPartRequest partRequest = partRequests.get(i); assertEquals(bucketName, partRequest.bucket()); @@ -1225,51 +1174,39 @@ public void testExecuteMultipartUploadIfEtagMatchesWithMetadataAndSSE() throws I } } - // Verify complete request CompleteMultipartUploadRequest completeRequest = completeRequestCaptor.getValue(); assertEquals(bucketName, completeRequest.bucket()); assertEquals(blobPath.buildAsString() + blobName, completeRequest.key()); assertEquals(uploadId, completeRequest.uploadId()); - // Verify ifMatch is on complete request - NOT create request assertEquals(inputETag, completeRequest.ifMatch()); - // Verify listener was called with exact final ETag verify(etagListener).onResponse(finalETag); verify(etagListener, never()).onFailure(any()); - // Verify resources were properly closed verify(clientReference).close(); } - /** - * Tests that actual content is properly transmitted during conditional multipart upload - */ public void testExecuteMultipartUploadIfEtagMatchesContentIntegrity() throws IOException { - // Setup test parameters final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); - final String inputETag = randomAlphaOfLengthBetween(8, 32); // Input ETag for condition - final String finalETag = randomAlphaOfLengthBetween(8, 32); // Final ETag returned + final String inputETag = randomAlphaOfLengthBetween(8, 32); + final String finalETag = randomAlphaOfLengthBetween(8, 32); final String uploadId = randomAlphaOfLengthBetween(10, 20); final BlobPath blobPath = new BlobPath(); - // Setup for multipart upload with known content final int partCount = 3; - final long partSize = ByteSizeUnit.MB.toBytes(5); // 5MB + final long partSize = ByteSizeUnit.MB.toBytes(5); final long blobSize = partSize * partCount; - // Create verifiable content - use a pattern that's different for each part final byte[] blobContent = new byte[(int) blobSize]; - Random random = new Random(0); // Fixed seed for reproducibility + Random random = new Random(0); random.nextBytes(blobContent); - // Mock S3BlobStore final S3BlobStore blobStore = mock(S3BlobStore.class); - // FIX 1: Replace mock StatsMetricPublisher with a real instance final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); when(blobStore.bucket()).thenReturn(bucketName); @@ -1281,13 +1218,11 @@ public void testExecuteMultipartUploadIfEtagMatchesContentIntegrity() throws IOE final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); - // Mock S3 client final S3Client client = mock(S3Client.class); final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); when(blobStore.clientReference()).thenReturn(clientReference); when(clientReference.get()).thenReturn(client); - // Setup argument captors with focus on request body final ArgumentCaptor createRequestCaptor = ArgumentCaptor.forClass( CreateMultipartUploadRequest.class ); @@ -1296,16 +1231,13 @@ public void testExecuteMultipartUploadIfEtagMatchesContentIntegrity() throws IOE ); final ArgumentCaptor requestBodyCaptor = ArgumentCaptor.forClass(RequestBody.class); - // Mock responses when(client.createMultipartUpload(createRequestCaptor.capture())).thenReturn( CreateMultipartUploadResponse.builder().uploadId(uploadId).build() ); - // Content verification will happen in this mock final List capturedPartContents = new ArrayList<>(); when(client.uploadPart(any(UploadPartRequest.class), requestBodyCaptor.capture())).thenAnswer(invocation -> { - // Capture content from request body for later verification RequestBody requestBody = requestBodyCaptor.getValue(); try (InputStream contentStream = requestBody.contentStreamProvider().newStream()) { byte[] partContent = contentStream.readAllBytes(); @@ -1314,21 +1246,16 @@ public void testExecuteMultipartUploadIfEtagMatchesContentIntegrity() throws IOE return UploadPartResponse.builder().eTag("etag-for-part").build(); }); - // Add final ETag to complete response when(client.completeMultipartUpload(completeRequestCaptor.capture())).thenReturn( CompleteMultipartUploadResponse.builder().eTag(finalETag).build() ); - // Create listener @SuppressWarnings("unchecked") ActionListener etagListener = mock(ActionListener.class); - // FIX 2: Handle exceptions properly - in this case we expect successful execution - // If we were expecting an exception, we would use expectThrows instead final ByteArrayInputStream inputStream = new ByteArrayInputStream(blobContent); blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, null, inputETag, etagListener); - // Verify key request parameters final CreateMultipartUploadRequest createRequest = createRequestCaptor.getValue(); assertEquals(bucketName, createRequest.bucket()); assertEquals(blobPath.buildAsString() + blobName, createRequest.key()); @@ -1356,9 +1283,6 @@ public void testExecuteMultipartUploadIfEtagMatchesContentIntegrity() throws IOE verify(clientReference).close(); } - /** - * Tests that multipart uploads enforce minimum and maximum size boundaries - */ public void testExecuteMultipartUploadIfEtagMatchesSizeValidation() { final S3BlobStore blobStore = mock(S3BlobStore.class); final S3BlobContainer blobContainer = new S3BlobContainer(mock(BlobPath.class), blobStore); @@ -1366,13 +1290,11 @@ public void testExecuteMultipartUploadIfEtagMatchesSizeValidation() { final String inputETag = randomAlphaOfLengthBetween(8, 32); final String finalETag = randomAlphaOfLengthBetween(8, 32); - // Use mocked listeners to verify interactions @SuppressWarnings("unchecked") ActionListener invalidSizeListener = mock(ActionListener.class); - // Test case 1: Below minimum size (4.9MB) { - final long tooSmallSize = ByteSizeUnit.MB.toBytes(5) - 1024; // Just below 5MB + final long tooSmallSize = ByteSizeUnit.MB.toBytes(5) - 1024; final IllegalArgumentException tooSmallException = expectThrows( IllegalArgumentException.class, @@ -1388,12 +1310,10 @@ public void testExecuteMultipartUploadIfEtagMatchesSizeValidation() { ); assertTrue(tooSmallException.getMessage().contains("can't be smaller than")); - // Listener should never be called when validation fails immediately verify(invalidSizeListener, never()).onResponse(any()); verify(invalidSizeListener, never()).onFailure(any()); } - // Test case 2: Above maximum size (5TB + 1 byte) { final long tooLargeSize = ByteSizeUnit.TB.toBytes(5) + 1; @@ -1415,7 +1335,6 @@ public void testExecuteMultipartUploadIfEtagMatchesSizeValidation() { verify(invalidSizeListener, never()).onFailure(any()); } - // Setup mocks for valid size tests final S3Client client = mock(S3Client.class); final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); @@ -1425,15 +1344,12 @@ public void testExecuteMultipartUploadIfEtagMatchesSizeValidation() { when(clientReference.get()).thenReturn(client); when(blobStore.bucket()).thenReturn("test-bucket"); - // Configure buffer size to allow valid tests when(blobStore.bufferSizeInBytes()).thenReturn(ByteSizeUnit.MB.toBytes(5)); - // Add argument captors for request validation ArgumentCaptor completeRequestCaptor = ArgumentCaptor.forClass( CompleteMultipartUploadRequest.class ); - // Setup success responses when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( CreateMultipartUploadResponse.builder().uploadId("test-upload-id").build() ); @@ -1446,13 +1362,11 @@ public void testExecuteMultipartUploadIfEtagMatchesSizeValidation() { CompleteMultipartUploadResponse.builder().eTag(finalETag).build() ); - // Test case 3: Exactly minimum size (5MB) { final long exactMinimumSize = ByteSizeUnit.MB.toBytes(5); @SuppressWarnings("unchecked") ActionListener validSizeListener = mock(ActionListener.class); - // Use a stream that returns zeros but doesn't allocate full memory InputStream zeroStream = new InputStream() { long remaining = exactMinimumSize; @@ -1479,7 +1393,6 @@ public int read(byte[] b, int off, int len) { }; try { - // This should not throw an exception blobContainer.executeMultipartUploadIfEtagMatches( blobStore, blobName, @@ -1490,14 +1403,11 @@ public int read(byte[] b, int off, int len) { validSizeListener ); - // Verify that listener was called with success verify(validSizeListener).onResponse(finalETag); verify(validSizeListener, never()).onFailure(any()); - // Verify resources were closed verify(clientReference).close(); - // Verify if-match condition was applied properly CompleteMultipartUploadRequest completeRequest = completeRequestCaptor.getValue(); assertEquals(inputETag, completeRequest.ifMatch()); @@ -1506,19 +1416,15 @@ public int read(byte[] b, int off, int len) { } } - // Reset for next test case reset(clientReference); when(blobStore.clientReference()).thenReturn(clientReference); when(clientReference.get()).thenReturn(client); - // Test case 4: Valid larger size { - // Use a moderate size for test purposes (10MB) final long testSize = ByteSizeUnit.MB.toBytes(10); @SuppressWarnings("unchecked") ActionListener validSizeListener = mock(ActionListener.class); - // Use a stream that returns zeros without allocating full memory InputStream zeroStream = new InputStream() { long remaining = testSize; @@ -1555,14 +1461,11 @@ public int read(byte[] b, int off, int len) { validSizeListener ); - // Verify that listener was called with success verify(validSizeListener).onResponse(finalETag); verify(validSizeListener, never()).onFailure(any()); - // Verify resources were closed verify(clientReference).close(); - // Verify if-match condition was applied properly CompleteMultipartUploadRequest completeRequest = completeRequestCaptor.getValue(); assertEquals(inputETag, completeRequest.ifMatch()); @@ -1572,20 +1475,15 @@ public int read(byte[] b, int off, int len) { } } - /** - * Tests handling of 412 Precondition Failed errors when ETag doesn't match - */ public void testExecuteMultipartUploadIfEtagMatchesPreconditionFailed() { - // Setup test parameters final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String eTag = randomAlphaOfLengthBetween(8, 32); final BlobPath blobPath = new BlobPath(); - final long partSize = ByteSizeUnit.MB.toBytes(5); // 5MB - final long blobSize = partSize * 2; // Two parts for simplicity + final long partSize = ByteSizeUnit.MB.toBytes(5); + final long blobSize = partSize * 2; - // Mock S3BlobStore final S3BlobStore blobStore = mock(S3BlobStore.class); final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); when(blobStore.bucket()).thenReturn(bucketName); @@ -1597,23 +1495,19 @@ public void testExecuteMultipartUploadIfEtagMatchesPreconditionFailed() { final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); - // Mock S3 client and references final S3Client client = mock(S3Client.class); final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); final AmazonS3Reference abortClientReference = mock(AmazonS3Reference.class); - // Setup client reference behavior when(blobStore.clientReference()).thenReturn(clientReference).thenReturn(abortClientReference); when(clientReference.get()).thenReturn(client); when(abortClientReference.get()).thenReturn(client); - // Create 412 Precondition Failed exception S3Exception preconditionFailedException = (S3Exception) S3Exception.builder() .message("Precondition Failed") - .statusCode(412) + .statusCode(S3BlobContainer.HTTP_STATUS_PRECONDITION_FAILED) .build(); - // Setup success for initial upload stages final String uploadId = randomAlphaOfLengthBetween(10, 20); when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( CreateMultipartUploadResponse.builder().uploadId(uploadId).build() @@ -1623,15 +1517,12 @@ public void testExecuteMultipartUploadIfEtagMatchesPreconditionFailed() { UploadPartResponse.builder().eTag("part-etag").build() ); - // Fail on complete with 412 when(client.completeMultipartUpload(any(CompleteMultipartUploadRequest.class))).thenThrow(preconditionFailedException); - // Setup abort response when(client.abortMultipartUpload(any(AbortMultipartUploadRequest.class))).thenReturn( AbortMultipartUploadResponse.builder().build() ); - // Execute the multipart upload final AtomicReference capturedException = new AtomicReference<>(); ActionListener etagListener = ActionListener.wrap( r -> fail("Should have failed with precondition failure"), @@ -1645,7 +1536,6 @@ public void testExecuteMultipartUploadIfEtagMatchesPreconditionFailed() { () -> blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, null, eTag, etagListener) ); - // Verify exception handling assertEquals("Unable to upload object [" + blobName + "] due to ETag mismatch", ioException.getMessage()); assertEquals(preconditionFailedException, ioException.getCause()); @@ -1654,21 +1544,15 @@ public void testExecuteMultipartUploadIfEtagMatchesPreconditionFailed() { assertTrue("Exception should be an OpenSearchException", exception instanceof OpenSearchException); assertEquals("stale_primary_shard", ((OpenSearchException) exception).getMessage()); - // Verify workflow verify(client).createMultipartUpload(any(CreateMultipartUploadRequest.class)); verify(client).completeMultipartUpload(any(CompleteMultipartUploadRequest.class)); verify(client).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); - // Verify both client references were closed verify(clientReference).close(); verify(abortClientReference).close(); } - /** - * Tests handling of various exception types during multipart uploads - */ public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { - // Setup basic test parameters final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String eTag = randomAlphaOfLengthBetween(8, 32); @@ -1677,30 +1561,23 @@ public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { final long partSize = ByteSizeUnit.MB.toBytes(5); final long blobSize = partSize * 2; - // Create test metadata Map metadata = new HashMap<>(); metadata.put("key1", "value1"); - // Define test cases with minimal necessary parameters - // [Exception type, error phase (0=create, 1=upload, 2=complete), status code] Object[][] testCases = { - // Key test cases for each error path - { "S3Exception", 0, 403 }, // Access denied during create - { "S3Exception", 1, 404 }, // Not found during upload - { "S3Exception", 2, 412 }, // Conditional check failure during complete - { "SdkException", 1, 0 }, // SDK error during upload - }; + { "S3Exception", 0, 403 }, + { "S3Exception", 1, 404 }, + { "S3Exception", 2, 412 }, + { "SdkException", 1, 0 }, }; for (Object[] testCase : testCases) { String exceptionType = (String) testCase[0]; int errorPhase = (int) testCase[1]; int statusCode = (int) testCase[2]; - // Create fresh mocks for each test case final S3BlobStore blobStore = mock(S3BlobStore.class); final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); - // Basic configuration when(blobStore.bucket()).thenReturn(bucketName); when(blobStore.bufferSizeInBytes()).thenReturn(partSize); when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); @@ -1709,17 +1586,14 @@ public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); - // Create client references final S3Client client = mock(S3Client.class); final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); final AmazonS3Reference abortClientReference = mock(AmazonS3Reference.class); - // Setup client reference behavior when(blobStore.clientReference()).thenReturn(clientReference).thenReturn(abortClientReference); when(clientReference.get()).thenReturn(client); when(abortClientReference.get()).thenReturn(client); - // Create the appropriate exception Exception testException; if ("S3Exception".equals(exceptionType)) { testException = (S3Exception) S3Exception.builder() @@ -1730,46 +1604,37 @@ public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { testException = SdkException.builder().message("SDK Error occurred").build(); } - // Setup responses for each phase final String uploadId = "test-upload-id"; if (errorPhase >= 1) { - // Success for create if error happens later when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( CreateMultipartUploadResponse.builder().uploadId(uploadId).build() ); } else { - // Error during create when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenThrow(testException); } if (errorPhase >= 2) { - // Success for upload parts if error happens at complete when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenReturn( UploadPartResponse.builder().eTag("part-etag").build() ); } else if (errorPhase == 1) { - // Error during upload when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenThrow(testException); } if (errorPhase == 2) { - // Error during complete when(client.completeMultipartUpload(any(CompleteMultipartUploadRequest.class))).thenThrow(testException); } - // Setup abort response when(client.abortMultipartUpload(any(AbortMultipartUploadRequest.class))).thenReturn( AbortMultipartUploadResponse.builder().build() ); - // Capture listener exception final AtomicReference capturedException = new AtomicReference<>(); ActionListener etagListener = ActionListener.wrap( r -> fail("Should have failed with exception"), capturedException::set ); - // Execute with expectThrows final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); IOException exception = expectThrows( IOException.class, @@ -1784,14 +1649,11 @@ public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { ) ); - // Verify thrown exception has the correct cause assertEquals(testException, exception.getCause()); - // Verify listener behavior Exception listenerException = capturedException.get(); assertNotNull("Expected a listener exception", listenerException); - // Special handling for 412 case if ("S3Exception".equals(exceptionType) && statusCode == 412) { assertTrue(listenerException instanceof OpenSearchException); assertEquals("stale_primary_shard", ((OpenSearchException) listenerException).getMessage()); @@ -1799,10 +1661,8 @@ public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { assertTrue(listenerException instanceof IOException); } - // Verify resource cleanup verify(clientReference).close(); - // Verify abort called when appropriate if (errorPhase >= 1) { verify(client).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); verify(abortClientReference).close(); @@ -1812,11 +1672,7 @@ public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { } } - /** - * Tests handling of generic SdkException scenarios during multipart uploads - */ public void testExecuteMultipartUploadIfEtagMatchesSdkException() { - // Test parameters final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String eTag = randomAlphaOfLengthBetween(8, 32); @@ -1825,19 +1681,16 @@ public void testExecuteMultipartUploadIfEtagMatchesSdkException() { final long partSize = ByteSizeUnit.MB.toBytes(5); final long blobSize = partSize * 2; - // Define test cases: [scenario name, error stage, should abort fail] Object[][] testScenarios = { - { "initialization error", 0, false }, // SdkException during createMultipartUpload - { "part upload error", 1, false }, // SdkException during uploadPart with successful abort - { "abort failure", 1, true } // SdkException during uploadPart with abort failure - }; + { "initialization error", 0, false }, + { "part upload error", 1, false }, + { "abort failure", 1, true } }; for (Object[] scenario : testScenarios) { String scenarioName = (String) scenario[0]; int errorStage = (int) scenario[1]; boolean abortFails = (boolean) scenario[2]; - // Setup fresh mocks for each scenario final S3BlobStore blobStore = mock(S3BlobStore.class); final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); when(blobStore.bucket()).thenReturn(bucketName); @@ -1847,7 +1700,6 @@ public void testExecuteMultipartUploadIfEtagMatchesSdkException() { final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); - // Mock S3 client final S3Client client = mock(S3Client.class); final AmazonS3Reference clientReference = mock(AmazonS3Reference.class); final AmazonS3Reference abortClientReference = mock(AmazonS3Reference.class); @@ -1855,25 +1707,19 @@ public void testExecuteMultipartUploadIfEtagMatchesSdkException() { when(clientReference.get()).thenReturn(client); when(abortClientReference.get()).thenReturn(client); - // Create primary exception SdkException primaryException = SdkException.builder().message("SDK error during " + scenarioName).build(); - // Setup responses based on scenario final String uploadId = "test-upload-id"; if (errorStage == 0) { - // Error during initialization when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenThrow(primaryException); } else { - // Success for initialization when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( CreateMultipartUploadResponse.builder().uploadId(uploadId).build() ); - // Error during part upload when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenThrow(primaryException); } - // Configure abort behavior with argument captor ArgumentCaptor abortCaptor = ArgumentCaptor.forClass(AbortMultipartUploadRequest.class); if (abortFails) { @@ -1883,14 +1729,12 @@ public void testExecuteMultipartUploadIfEtagMatchesSdkException() { when(client.abortMultipartUpload(abortCaptor.capture())).thenReturn(AbortMultipartUploadResponse.builder().build()); } - // Capture listener exception final AtomicReference capturedException = new AtomicReference<>(); ActionListener etagListener = ActionListener.wrap( r -> fail("Should have failed with SdkException"), capturedException::set ); - // Execute with expectThrows final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); IOException exception = expectThrows( IOException.class, @@ -1905,24 +1749,18 @@ public void testExecuteMultipartUploadIfEtagMatchesSdkException() { ) ); - // Verify thrown exception assertEquals("Original SdkException should be preserved as cause", primaryException, exception.getCause()); - // Verify listener exception Exception listenerException = capturedException.get(); assertNotNull("Expected an exception", listenerException); assertTrue("Exception should be IOException", listenerException instanceof IOException); - // Verify main client reference was closed verify(clientReference).close(); - // Verify abort behavior if (errorStage > 0) { - // Should attempt abort if upload ID exists verify(client).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); verify(abortClientReference).close(); - // Verify abort request parameters if (!abortCaptor.getAllValues().isEmpty()) { AbortMultipartUploadRequest abortRequest = abortCaptor.getValue(); assertEquals("Abort request should have correct upload ID", uploadId, abortRequest.uploadId()); @@ -1930,17 +1768,12 @@ public void testExecuteMultipartUploadIfEtagMatchesSdkException() { assertEquals("Abort request should have correct key", blobPath.buildAsString() + blobName, abortRequest.key()); } } else { - // No abort if no upload ID verify(client, never()).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); } } } - /** - * Tests that client references are properly managed (acquired and closed) in various scenarios - */ public void testExecuteMultipartUploadIfEtagMatchesResourceManagement() throws IOException { - // Setup test parameters final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String eTag = randomAlphaOfLengthBetween(8, 32); @@ -1950,17 +1783,15 @@ public void testExecuteMultipartUploadIfEtagMatchesResourceManagement() throws I final long partSize = ByteSizeUnit.MB.toBytes(5); final long blobSize = partSize * 2; - // Define scenarios to test different resource management cases enum ResourceScenario { - SUCCESS_PATH, // Complete successful upload - INIT_FAILURE, // Failure during createMultipartUpload - PART_UPLOAD_FAILURE, // Failure during uploadPart - COMPLETION_FAILURE, // Failure during completeMultipartUpload - ABORT_FAILURE // Failure during abortMultipartUpload + SUCCESS_PATH, + INIT_FAILURE, + PART_UPLOAD_FAILURE, + COMPLETION_FAILURE, + ABORT_FAILURE } for (ResourceScenario scenario : ResourceScenario.values()) { - // Create fresh mocks for each scenario final S3BlobStore blobStore = mock(S3BlobStore.class); final StatsMetricPublisher metricPublisher = new StatsMetricPublisher(); when(blobStore.bucket()).thenReturn(bucketName); @@ -1968,31 +1799,24 @@ enum ResourceScenario { when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); - // S3 client and references final S3Client client = mock(S3Client.class); - // The implementation always creates separate references for the main operation and abort AmazonS3Reference primaryClientReference = mock(AmazonS3Reference.class); AmazonS3Reference abortClientReference = mock(AmazonS3Reference.class); when(primaryClientReference.get()).thenReturn(client); when(abortClientReference.get()).thenReturn(client); - // Configure clientReference behavior based on implementation pattern when(blobStore.clientReference()).thenReturn(primaryClientReference).thenReturn(abortClientReference); final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); - // Create exceptions with distinctive messages for each scenario SdkException stageException = SdkException.builder().message("Failure during " + scenario.name()).build(); - // Setup argument captor for abort request validation ArgumentCaptor abortCaptor = ArgumentCaptor.forClass(AbortMultipartUploadRequest.class); - // Configure upload behavior based on scenario switch (scenario) { case SUCCESS_PATH: - // Successful path should close reference after completion when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( CreateMultipartUploadResponse.builder().uploadId(uploadId).build() ); @@ -2005,12 +1829,10 @@ enum ResourceScenario { break; case INIT_FAILURE: - // Initialization failure should close reference immediately when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenThrow(stageException); break; case PART_UPLOAD_FAILURE: - // Part upload failure should close primary reference and abort reference when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( CreateMultipartUploadResponse.builder().uploadId(uploadId).build() ); @@ -2019,7 +1841,6 @@ enum ResourceScenario { break; case COMPLETION_FAILURE: - // Completion failure should close primary reference and abort reference when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( CreateMultipartUploadResponse.builder().uploadId(uploadId).build() ); @@ -2031,7 +1852,6 @@ enum ResourceScenario { break; case ABORT_FAILURE: - // Even if abort fails, all references should be closed when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( CreateMultipartUploadResponse.builder().uploadId(uploadId).build() ); @@ -2042,26 +1862,20 @@ enum ResourceScenario { break; } - // Create action listener to capture callbacks @SuppressWarnings("unchecked") ActionListener etagListener = mock(ActionListener.class); final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); - // Execute the appropriate test based on scenario if (scenario == ResourceScenario.SUCCESS_PATH) { - // Success path should not throw exceptions blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, null, eTag, etagListener); - // Verify listener success callback verify(etagListener).onResponse("final-etag"); verify(etagListener, never()).onFailure(any()); - // Success path uses one client reference exactly once verify(blobStore, times(1)).clientReference(); verify(primaryClientReference).close(); } else { - // For all failure scenarios, expect exception IOException exception = expectThrows( IOException.class, () -> blobContainer.executeMultipartUploadIfEtagMatches( @@ -2075,23 +1889,18 @@ enum ResourceScenario { ) ); - // Verify exception has the correct cause assertEquals("Exception cause should be the original exception", stageException, exception.getCause()); - // Verify listener failure callback verify(etagListener).onFailure(any(Exception.class)); verify(etagListener, never()).onResponse(any()); - // Primary reference should always be closed verify(primaryClientReference).close(); if (scenario != ResourceScenario.INIT_FAILURE) { - // For non-init failures, verify abort is called - verify(blobStore, times(2)).clientReference(); // Main + abort + verify(blobStore, times(2)).clientReference(); verify(client).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); verify(abortClientReference).close(); - // Verify abort request parameters if (!abortCaptor.getAllValues().isEmpty()) { AbortMultipartUploadRequest abortRequest = abortCaptor.getValue(); assertEquals("Upload ID should match", uploadId, abortRequest.uploadId()); @@ -2099,8 +1908,7 @@ enum ResourceScenario { assertEquals("Key should match", blobPath.buildAsString() + blobName, abortRequest.key()); } } else { - // Init failure should not use abort - verify(blobStore, times(1)).clientReference(); // Main only + verify(blobStore, times(1)).clientReference(); verify(client, never()).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); } } From 44d45e86e69dba4ab184ffb53c94b15765a684c3 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Tue, 27 May 2025 16:51:09 +0530 Subject: [PATCH 04/20] UploadRequest : Add ConditionalWriteOptions --- .../repositories/s3/async/UploadRequest.java | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/UploadRequest.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/UploadRequest.java index 79b58ff215c54..ea2dffe2b6039 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/UploadRequest.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/UploadRequest.java @@ -10,6 +10,7 @@ import org.opensearch.common.CheckedConsumer; import org.opensearch.common.Nullable; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteOptions; import org.opensearch.common.blobstore.stream.write.WritePriority; import java.io.IOException; @@ -28,6 +29,7 @@ public class UploadRequest { private final Long expectedChecksum; private final Map metadata; private final boolean uploadRetryEnabled; + private final ConditionalWriteOptions conditionalOptions; /** * Construct a new UploadRequest object @@ -40,6 +42,7 @@ public class UploadRequest { * @param doRemoteDataIntegrityCheck A boolean to inform vendor plugins whether remote data integrity checks need to be done * @param expectedChecksum Checksum of the file being uploaded for remote data integrity check * @param metadata Metadata of the file being uploaded + * @param conditionalOptions Conditions that must be satisfied for the write to succeed */ public UploadRequest( String bucket, @@ -50,7 +53,8 @@ public UploadRequest( boolean doRemoteDataIntegrityCheck, Long expectedChecksum, boolean uploadRetryEnabled, - @Nullable Map metadata + @Nullable Map metadata, + @Nullable ConditionalWriteOptions conditionalOptions ) { this.bucket = bucket; this.key = key; @@ -61,6 +65,7 @@ public UploadRequest( this.expectedChecksum = expectedChecksum; this.uploadRetryEnabled = uploadRetryEnabled; this.metadata = metadata; + this.conditionalOptions = conditionalOptions; } public String getBucket() { @@ -101,4 +106,12 @@ public boolean isUploadRetryEnabled() { public Map getMetadata() { return metadata; } + + /** + * @return conditional write options for this upload, or null if none are specified + */ + @Nullable + public ConditionalWriteOptions getConditionalOptions() { + return conditionalOptions; + } } From a31d15e53f86c4cfedfdd07d2b40e63d2b7c4afc Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Tue, 27 May 2025 16:52:15 +0530 Subject: [PATCH 05/20] Add ConditionalWrite utility for conditional blob store writes Introduces ConditionalWriteOptions and ConditionalWriteResponse classes to support conditional write operations in BlobContainer implementations. --- .../common/blobstore/ConditionalWrite.java | 167 ++++++++++++++++++ 1 file changed, 167 insertions(+) create mode 100644 server/src/main/java/org/opensearch/common/blobstore/ConditionalWrite.java diff --git a/server/src/main/java/org/opensearch/common/blobstore/ConditionalWrite.java b/server/src/main/java/org/opensearch/common/blobstore/ConditionalWrite.java new file mode 100644 index 0000000000000..0b3f94aa603b8 --- /dev/null +++ b/server/src/main/java/org/opensearch/common/blobstore/ConditionalWrite.java @@ -0,0 +1,167 @@ +/* + * 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.common.blobstore; + +import java.time.Instant; + +/** + * Utility classes supporting conditional write operations on a {@link BlobContainer}. + * The main entry points are {@link ConditionalWriteOptions} for specifying conditions, and + * {@link ConditionalWriteResponse} for receiving the result of a conditional write. + */ +public final class ConditionalWrite { + private ConditionalWrite() {} + + /** + * Encapsulates options controlling preconditions to be deployed when a blob is to be written to the remote store. + * Immutable and thread-safe. Use the provided static factory methods or the {@link Builder} + * to construct instances with the desired conditions. These options can be supplied to + * blob store write operations to enforce preconditions + * + */ + public static final class ConditionalWriteOptions { + + private final boolean ifNotExists; + private final boolean ifMatch; + private final boolean ifUnmodifiedSince; + private final String versionIdentifier; + private final Instant unmodifiedSince; + + private ConditionalWriteOptions(Builder builder) { + this.ifNotExists = builder.ifNotExists; + this.ifMatch = builder.ifMatch; + this.ifUnmodifiedSince = builder.ifUnmodifiedSince; + this.versionIdentifier = builder.versionIdentifier; + this.unmodifiedSince = builder.unmodifiedSince; + } + + public static ConditionalWriteOptions none() { + return new Builder().build(); + } + + public static ConditionalWriteOptions ifNotExists() { + return new Builder().setIfNotExists(true).build(); + } + + public static ConditionalWriteOptions ifMatch(String versionIdentifier) { + return new Builder().setIfMatch(true).setVersionIdentifier(versionIdentifier).build(); + } + + public static ConditionalWriteOptions ifUnmodifiedSince(Instant ts) { + return new Builder().setIfUnmodifiedSince(true).setUnmodifiedSince(ts).build(); + } + + /** + * Returns a new {@link Builder} for constructing custom conditional write options. + */ + public static Builder builder() { + return new Builder(); + } + + public boolean isIfNotExists() { + return ifNotExists; + } + + public boolean isIfMatch() { + return ifMatch; + } + + public boolean isIfUnmodifiedSince() { + return ifUnmodifiedSince; + } + + public String getVersionIdentifier() { + return versionIdentifier; + } + + public Instant getUnmodifiedSince() { + return unmodifiedSince; + } + + public static final class Builder { + private boolean ifNotExists = false; + private boolean ifMatch = false; + private boolean ifUnmodifiedSince = false; + private String versionIdentifier = null; + private Instant unmodifiedSince = null; + + private Builder() {} + + /** + * Sets the write to succeed only if the blob does not exist. + * @param flag true to enable this condition + * @return this builder + */ + public Builder setIfNotExists(boolean flag) { + this.ifNotExists = flag; + return this; + } + + /** + * Sets the write to succeed only if the blob matches the expected version. + * @param flag true to enable this condition + * @return this builder + */ + public Builder setIfMatch(boolean flag) { + this.ifMatch = flag; + return this; + } + + /** + * Sets the write to succeed only if the blob was not modified since a given instant. + * @param flag true to enable this condition + * @return this builder + */ + public Builder setIfUnmodifiedSince(boolean flag) { + this.ifUnmodifiedSince = flag; + return this; + } + + /** + * Sets the timestamp before which the blob must remain unmodified. + * @param ts the instant to check + * @return this builder + */ + public Builder setUnmodifiedSince(Instant ts) { + this.unmodifiedSince = ts; + return this; + } + + public Builder setVersionIdentifier(String versionIdentifier) { + this.versionIdentifier = versionIdentifier; + return this; + } + + public ConditionalWriteOptions build() { + return new ConditionalWriteOptions(this); + } + } + } + + /** + * encapsulates the result of a conditional write operation. + * Contains the new version identifier (such as an ETag or version string) retrieved from the remote store + * after a successful write. + */ + public static final class ConditionalWriteResponse { + private final String newVersionIdentifier; + + private ConditionalWriteResponse(String versionIdentifier) { + this.newVersionIdentifier = versionIdentifier; + } + + public static ConditionalWriteResponse success(String versionIdentifier) { + return new ConditionalWriteResponse(versionIdentifier); + } + + public String getVersionIdentifier() { + return newVersionIdentifier; + } + } +} From a1df80af9e6d521b06a4e72589e00774bfd17a59 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Wed, 28 May 2025 02:30:01 +0530 Subject: [PATCH 06/20] Adds logic to Upload an object to S3 conditionally using the async client --- .../s3/async/AsyncTransferManager.java | 370 ++++++++++++++++++ 1 file changed, 370 insertions(+) diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java index 0f9bf3be77d73..143417d238454 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java @@ -31,6 +31,8 @@ import org.apache.logging.log4j.message.ParameterizedMessage; import org.opensearch.ExceptionsHelper; import org.opensearch.common.StreamContext; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteOptions; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteResponse; import org.opensearch.common.blobstore.exception.CorruptFileException; import org.opensearch.common.blobstore.stream.write.WritePriority; import org.opensearch.common.io.InputStreamContainer; @@ -68,6 +70,9 @@ public final class AsyncTransferManager { private final long minimumPartSize; private final long maxRetryablePartSize; + private static final int HTTP_STATUS_PRECONDITION_FAILED = 412; + private static final int HTTP_STATUS_CONFLICT = 409; + @SuppressWarnings("rawtypes") private final TransferSemaphoresHolder transferSemaphoresHolder; @@ -143,6 +148,63 @@ public CompletableFuture uploadObject( return returnFuture; } + /** + * Upload an object to S3 conditionally using the async client + * + * @param s3AsyncClient S3 client to use for upload + * @param uploadRequest The {@link UploadRequest} object encapsulating all relevant details for upload + * @param streamContext The {@link StreamContext} to supply streams during upload + * @param statsMetricPublisher Metric publisher for collecting stats + * @return A {@link CompletableFuture} that will complete with the ConditionalWriteResponse or an exception + */ + public CompletableFuture uploadObjectConditionally( + S3AsyncClient s3AsyncClient, + UploadRequest uploadRequest, + StreamContext streamContext, + StatsMetricPublisher statsMetricPublisher + ) { + ConditionalWriteOptions options = uploadRequest.getConditionalOptions(); + if (options == null) { + throw new IllegalArgumentException("Cannot perform conditional upload with null options"); + } + + CompletableFuture returnFuture = new CompletableFuture<>(); + try { + if (streamContext.getNumberOfParts() == 1) { + log.debug(() -> "Starting conditional single part upload for key: " + uploadRequest.getKey()); + TransferSemaphoresHolder.RequestContext requestContext = transferSemaphoresHolder.createRequestContext(); + Semaphore semaphore = AsyncPartsHandler.maybeAcquireSemaphore( + transferSemaphoresHolder, + requestContext, + uploadRequest.getWritePriority(), + uploadRequest.getKey() + ); + try { + uploadInOneChunkConditionally( + s3AsyncClient, + uploadRequest, + streamContext, + returnFuture, + statsMetricPublisher, + semaphore + ); + } catch (Exception ex) { + if (semaphore != null) { + semaphore.release(); + } + throw ex; + } + } else { + log.debug(() -> "Starting conditional multipart upload for key: " + uploadRequest.getKey()); + uploadInPartsConditionally(s3AsyncClient, uploadRequest, streamContext, returnFuture, statsMetricPublisher); + } + } catch (Throwable throwable) { + returnFuture.completeExceptionally(throwable); + } + + return returnFuture; + } + private void uploadInParts( S3AsyncClient s3AsyncClient, UploadRequest uploadRequest, @@ -184,6 +246,47 @@ private void uploadInParts( doUploadInParts(s3AsyncClient, uploadRequest, streamContext, returnFuture, uploadId, statsMetricPublisher); } + private void uploadInPartsConditionally( + S3AsyncClient s3AsyncClient, + UploadRequest uploadRequest, + StreamContext streamContext, + CompletableFuture returnFuture, + StatsMetricPublisher statsMetricPublisher + ) { + CreateMultipartUploadRequest.Builder createMultipartUploadRequestBuilder = CreateMultipartUploadRequest.builder() + .bucket(uploadRequest.getBucket()) + .key(uploadRequest.getKey()) + .overrideConfiguration(o -> o.addMetricPublisher(statsMetricPublisher.multipartUploadMetricCollector)); + + if (CollectionUtils.isNotEmpty(uploadRequest.getMetadata())) { + createMultipartUploadRequestBuilder.metadata(uploadRequest.getMetadata()); + } + if (uploadRequest.doRemoteDataIntegrityCheck()) { + createMultipartUploadRequestBuilder.checksumAlgorithm(ChecksumAlgorithm.CRC32); + } + + CompletableFuture createMultipartUploadFuture = SocketAccess.doPrivileged( + () -> s3AsyncClient.createMultipartUpload(createMultipartUploadRequestBuilder.build()) + ); + + // Ensure cancellations are forwarded to the createMultipartUploadFuture future + CompletableFutureUtils.forwardExceptionTo(returnFuture, createMultipartUploadFuture); + + String uploadId; + try { + // Block main thread here so that upload of parts doesn't get executed in future completion thread. + // We should never execute latent operation like acquisition of permit in future completion pool. + CreateMultipartUploadResponse createMultipartUploadResponse = createMultipartUploadFuture.get(); + uploadId = createMultipartUploadResponse.uploadId(); + log.debug(() -> "Initiated new multipart upload, uploadId: " + createMultipartUploadResponse.uploadId()); + } catch (Exception ex) { + handleConditionalException(returnFuture, ex, uploadRequest.getKey()); + return; + } + + doUploadInPartsConditionally(s3AsyncClient, uploadRequest, streamContext, returnFuture, uploadId, statsMetricPublisher); + } + private void doUploadInParts( S3AsyncClient s3AsyncClient, UploadRequest uploadRequest, @@ -244,6 +347,112 @@ private void doUploadInParts( }); } + private void doUploadInPartsConditionally( + S3AsyncClient s3AsyncClient, + UploadRequest uploadRequest, + StreamContext streamContext, + CompletableFuture returnFuture, + String uploadId, + StatsMetricPublisher statsMetricPublisher + ) { + // The list of completed parts must be sorted + AtomicReferenceArray completedParts = new AtomicReferenceArray<>(streamContext.getNumberOfParts()); + AtomicReferenceArray inputStreamContainers = new AtomicReferenceArray<>(streamContext.getNumberOfParts()); + + List> futures; + try { + futures = AsyncPartsHandler.uploadParts( + s3AsyncClient, + executorService, + priorityExecutorService, + urgentExecutorService, + uploadRequest, + streamContext, + uploadId, + completedParts, + inputStreamContainers, + statsMetricPublisher, + uploadRequest.isUploadRetryEnabled(), + transferSemaphoresHolder, + maxRetryablePartSize + ); + } catch (Exception ex) { + try { + AsyncPartsHandler.cleanUpParts(s3AsyncClient, uploadRequest, uploadId); + } finally { + returnFuture.completeExceptionally(ex); + } + return; + } + + CompletableFutureUtils.allOfExceptionForwarded(futures.toArray(CompletableFuture[]::new)).thenApply(resp -> { + try { + uploadRequest.getUploadFinalizer().accept(true); + } catch (IOException e) { + throw new RuntimeException(e); + } + return resp; + }).thenApply(ignore -> { + if (uploadRequest.doRemoteDataIntegrityCheck()) { + mergeAndVerifyChecksum(inputStreamContainers, uploadRequest.getKey(), uploadRequest.getExpectedChecksum()); + } + return null; + }).thenCompose(ignore -> { + log.debug(() -> "Completing conditional multipart upload, uploadId: " + uploadId); + + CompletedPart[] parts = IntStream.range(0, completedParts.length()).mapToObj(completedParts::get).toArray(CompletedPart[]::new); + + CompleteMultipartUploadRequest.Builder completeRequestBuilder = CompleteMultipartUploadRequest.builder() + .bucket(uploadRequest.getBucket()) + .key(uploadRequest.getKey()) + .uploadId(uploadId) + .multipartUpload(CompletedMultipartUpload.builder().parts(parts).build()) + .overrideConfiguration(o -> o.addMetricPublisher(statsMetricPublisher.multipartUploadMetricCollector)); + + if (uploadRequest.getConditionalOptions() != null) { + applyConditionalHeaders(completeRequestBuilder, uploadRequest.getConditionalOptions()); + } + + return s3AsyncClient.completeMultipartUpload(completeRequestBuilder.build()); + }).handle((response, throwable) -> { + if (throwable != null) { + AsyncPartsHandler.cleanUpParts(s3AsyncClient, uploadRequest, uploadId); + + Throwable unwrappedThrowable = ExceptionsHelper.unwrap(throwable, S3Exception.class); + if (unwrappedThrowable != null) { + S3Exception s3Exception = (S3Exception) unwrappedThrowable; + if (s3Exception.statusCode() == HTTP_STATUS_PRECONDITION_FAILED) { + returnFuture.completeExceptionally( + S3Exception.builder() + .message("Conditional write failed: condition not met for " + uploadRequest.getKey()) + .statusCode(HTTP_STATUS_PRECONDITION_FAILED) + .cause(s3Exception) + .build() + ); + return null; + } else if (s3Exception.statusCode() == HTTP_STATUS_CONFLICT) { + returnFuture.completeExceptionally( + S3Exception.builder() + .message("Blob already exists: " + uploadRequest.getKey()) + .statusCode(HTTP_STATUS_CONFLICT) + .cause(s3Exception) + .build() + ); + return null; + } + } + handleConditionalException(returnFuture, throwable, uploadRequest.getKey()); + return null; + } else { + returnFuture.complete(ConditionalWriteResponse.success(response.eTag())); + return null; + } + }).exceptionally(throwable -> { + handleConditionalException(returnFuture, throwable, uploadRequest.getKey()); + return null; + }); + } + private void mergeAndVerifyChecksum( AtomicReferenceArray inputStreamContainers, String fileName, @@ -328,6 +537,44 @@ private static void handleException(CompletableFuture returnFuture, Suppli } } + /** + * Error handler for conditional uploads + */ + private void handleConditionalException(CompletableFuture returnFuture, Throwable throwable, String resourceName) { + Throwable cause = throwable; + while (cause instanceof CompletionException && cause.getCause() != null) { + cause = cause.getCause(); + } + + if (cause instanceof S3Exception s3e) { + if (s3e.statusCode() == HTTP_STATUS_PRECONDITION_FAILED) { + returnFuture.completeExceptionally( + S3Exception.builder() + .message("Conditional write failed: condition not met for " + resourceName) + .statusCode(HTTP_STATUS_PRECONDITION_FAILED) + .cause(s3e) + .build() + ); + return; + } else if (s3e.statusCode() == HTTP_STATUS_CONFLICT) { + returnFuture.completeExceptionally( + S3Exception.builder() + .message("Blob already exists: " + resourceName) + .statusCode(HTTP_STATUS_CONFLICT) + .cause(s3e) + .build() + ); + return; + } + } + if (cause instanceof Error) { + returnFuture.completeExceptionally(cause); + } else { + SdkClientException exception = SdkClientException.create("Failed conditional upload of " + resourceName, cause); + returnFuture.completeExceptionally(exception); + } + } + /** * Calculates the optimal part size of each part request if the upload operation is carried out as multipart upload. */ @@ -436,6 +683,129 @@ private void uploadInOneChunk( CompletableFutureUtils.forwardResultTo(putObjectFuture, returnFuture); } + private void uploadInOneChunkConditionally( + S3AsyncClient s3AsyncClient, + UploadRequest uploadRequest, + StreamContext streamContext, + CompletableFuture returnFuture, + StatsMetricPublisher statsMetricPublisher, + Semaphore semaphore + ) { + PutObjectRequest.Builder putObjectRequestBuilder = PutObjectRequest.builder() + .bucket(uploadRequest.getBucket()) + .key(uploadRequest.getKey()) + .contentLength(uploadRequest.getContentLength()) + .overrideConfiguration(o -> o.addMetricPublisher(statsMetricPublisher.putObjectMetricPublisher)); + + if (CollectionUtils.isNotEmpty(uploadRequest.getMetadata())) { + putObjectRequestBuilder.metadata(uploadRequest.getMetadata()); + } + if (uploadRequest.doRemoteDataIntegrityCheck()) { + putObjectRequestBuilder.checksumAlgorithm(ChecksumAlgorithm.CRC32); + putObjectRequestBuilder.checksumCRC32(base64StringFromLong(uploadRequest.getExpectedChecksum())); + } + + if (uploadRequest.getConditionalOptions() != null) { + applyConditionalHeaders(putObjectRequestBuilder, uploadRequest.getConditionalOptions()); + } + + PutObjectRequest putObjectRequest = putObjectRequestBuilder.build(); + ExecutorService streamReadExecutor; + if (uploadRequest.getWritePriority() == WritePriority.URGENT) { + streamReadExecutor = urgentExecutorService; + } else if (uploadRequest.getWritePriority() == WritePriority.HIGH) { + streamReadExecutor = priorityExecutorService; + } else { + streamReadExecutor = executorService; + } + + CompletableFuture putObjectFuture = SocketAccess.doPrivileged(() -> { + InputStream inputStream = null; + CompletableFuture putObjectRespFuture; + try { + InputStreamContainer inputStreamContainer = streamContext.provideStream(0); + inputStream = AsyncPartsHandler.maybeRetryInputStream( + inputStreamContainer.getInputStream(), + uploadRequest.getWritePriority(), + uploadRequest.isUploadRetryEnabled(), + uploadRequest.getContentLength(), + maxRetryablePartSize + ); + AsyncRequestBody asyncRequestBody = AsyncRequestBody.fromInputStream( + inputStream, + inputStreamContainer.getContentLength(), + streamReadExecutor + ); + putObjectRespFuture = s3AsyncClient.putObject(putObjectRequest, asyncRequestBody); + } catch (Exception e) { + releaseResourcesSafely(semaphore, inputStream, uploadRequest.getKey()); + return CompletableFuture.failedFuture(e); + } + + InputStream finalInputStream = inputStream; + + return putObjectRespFuture.handle((resp, throwable) -> { + releaseResourcesSafely(semaphore, finalInputStream, uploadRequest.getKey()); + + if (throwable != null) { + Throwable unwrappedThrowable = ExceptionsHelper.unwrap(throwable, S3Exception.class); + if (unwrappedThrowable != null) { + S3Exception s3Exception = (S3Exception) unwrappedThrowable; + if (s3Exception.statusCode() == HttpStatusCode.BAD_REQUEST + && "BadDigest".equals(s3Exception.awsErrorDetails().errorCode())) { + throw new RuntimeException(new CorruptFileException(s3Exception, uploadRequest.getKey())); + } + } + returnFuture.completeExceptionally(throwable); + } else { + try { + uploadRequest.getUploadFinalizer().accept(true); + } catch (IOException e) { + throw new RuntimeException(e); + } + returnFuture.complete(ConditionalWriteResponse.success(resp.eTag())); + } + + return null; + }).handle((resp, throwable) -> { + if (throwable != null) { + deleteUploadedObject(s3AsyncClient, uploadRequest); + returnFuture.completeExceptionally(throwable); + } + + return null; + }); + }); + + CompletableFutureUtils.forwardExceptionTo(returnFuture, putObjectFuture); + CompletableFutureUtils.forwardResultTo(putObjectFuture, returnFuture); + } + + /** + * Apply conditional headers to a request builder + */ + private void applyConditionalHeaders(Object builder, ConditionalWriteOptions options) { + if (options == null) { + return; + } + + if (builder instanceof PutObjectRequest.Builder) { + PutObjectRequest.Builder putBuilder = (PutObjectRequest.Builder) builder; + if (options.isIfNotExists()) { + putBuilder.ifNoneMatch("*"); + } else if (options.isIfMatch()) { + putBuilder.ifMatch(options.getVersionIdentifier()); + } + } else if (builder instanceof CompleteMultipartUploadRequest.Builder) { + CompleteMultipartUploadRequest.Builder completeBuilder = (CompleteMultipartUploadRequest.Builder) builder; + if (options.isIfNotExists()) { + completeBuilder.ifNoneMatch("*"); + } else if (options.isIfMatch()) { + completeBuilder.ifMatch(options.getVersionIdentifier()); + } + } + } + private void releaseResourcesSafely(Semaphore semaphore, InputStream inputStream, String file) { if (semaphore != null) { semaphore.release(); From 2c2c72a5e834596afdd20a28bf1a7c16f08c0709 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Wed, 28 May 2025 02:30:27 +0530 Subject: [PATCH 07/20] UTs for async Conditional upload logic --- .../s3/async/AsyncTransferManagerTests.java | 331 +++++++++++++++++- 1 file changed, 327 insertions(+), 4 deletions(-) diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java index 89add3cdbfc60..36321f80e7f7d 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java @@ -28,6 +28,8 @@ import org.opensearch.ExceptionsHelper; import org.opensearch.common.StreamContext; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteOptions; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteResponse; import org.opensearch.common.blobstore.exception.CorruptFileException; import org.opensearch.common.blobstore.stream.write.WritePriority; import org.opensearch.common.io.InputStreamContainer; @@ -97,7 +99,7 @@ public void testOneChunkUpload() { s3AsyncClient, new UploadRequest("bucket", "key", ByteSizeUnit.MB.toBytes(1), WritePriority.HIGH, uploadSuccess -> { // do nothing - }, false, null, true, metadata), + }, false, null, true, metadata, null), new StreamContext((partIdx, partSize, position) -> { streamRef.set(new ZeroInputStream(partSize)); return new InputStreamContainer(streamRef.get(), partSize, position); @@ -146,7 +148,7 @@ public void testOneChunkUploadCorruption() { s3AsyncClient, new UploadRequest("bucket", "key", ByteSizeUnit.MB.toBytes(1), WritePriority.HIGH, uploadSuccess -> { // do nothing - }, false, null, true, metadata), + }, false, null, true, metadata, null), new StreamContext( (partIdx, partSize, position) -> new InputStreamContainer(new ZeroInputStream(partSize), partSize, position), ByteSizeUnit.MB.toBytes(1), @@ -203,7 +205,7 @@ public void testMultipartUpload() { s3AsyncClient, new UploadRequest("bucket", "key", ByteSizeUnit.MB.toBytes(5), WritePriority.HIGH, uploadSuccess -> { // do nothing - }, true, 3376132981L, true, metadata), + }, true, 3376132981L, true, metadata, null), new StreamContext((partIdx, partSize, position) -> { InputStream stream = new ZeroInputStream(partSize); streams.add(stream); @@ -267,7 +269,7 @@ public void testMultipartUploadCorruption() { s3AsyncClient, new UploadRequest("bucket", "key", ByteSizeUnit.MB.toBytes(5), WritePriority.HIGH, uploadSuccess -> { // do nothing - }, true, 0L, true, metadata), + }, true, 0L, true, metadata, null), new StreamContext( (partIdx, partSize, position) -> new InputStreamContainer(new ZeroInputStream(partSize), partSize, position), ByteSizeUnit.MB.toBytes(1), @@ -291,4 +293,325 @@ public void testMultipartUploadCorruption() { verify(s3AsyncClient, times(0)).completeMultipartUpload(any(CompleteMultipartUploadRequest.class)); verify(s3AsyncClient, times(1)).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); } + + public void testConditionalOneChunkUpload() { + CompletableFuture putObjectResponseCompletableFuture = new CompletableFuture<>(); + putObjectResponseCompletableFuture.complete(PutObjectResponse.builder().eTag("test-etag-1234").build()); + when(s3AsyncClient.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))).thenReturn( + putObjectResponseCompletableFuture + ); + + ConditionalWriteOptions options = ConditionalWriteOptions.ifMatch("old-etag-value"); + AtomicReference streamRef = new AtomicReference<>(); + Map metadata = new HashMap<>(); + metadata.put("key1", "value1"); + metadata.put("key2", "value2"); + + CompletableFuture resultFuture = asyncTransferManager.uploadObjectConditionally( + s3AsyncClient, + new UploadRequest( + "bucket", + "key", + ByteSizeUnit.MB.toBytes(1), + WritePriority.HIGH, + uploadSuccess -> {}, + false, + null, + true, + metadata, + options + ), + new StreamContext((partIdx, partSize, position) -> { + streamRef.set(new ZeroInputStream(partSize)); + return new InputStreamContainer(streamRef.get(), partSize, position); + }, ByteSizeUnit.MB.toBytes(1), ByteSizeUnit.MB.toBytes(1), 1), + new StatsMetricPublisher() + ); + + try { + ConditionalWriteResponse response = resultFuture.get(); + assertNotNull("Response should not be null", response); + assertEquals("ETag should match expected value", "test-etag-1234", response.getVersionIdentifier()); + } catch (ExecutionException | InterruptedException e) { + fail("Did not expect resultFuture to fail: " + e.getMessage()); + } + + verify(s3AsyncClient, times(1)).putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class)); + + boolean closeError = false; + try { + streamRef.get().available(); + } catch (IOException e) { + closeError = e.getMessage().equals("Stream closed"); + } + assertTrue("InputStream was still open after upload", closeError); + } + + public void testConditionalOneChunkUploadPreconditionFailed() { + CompletableFuture putObjectResponseCompletableFuture = new CompletableFuture<>(); + S3Exception mockException = (S3Exception) S3Exception.builder().statusCode(412).message("Precondition Failed").build(); + + putObjectResponseCompletableFuture.completeExceptionally(mockException); + when(s3AsyncClient.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))).thenReturn( + putObjectResponseCompletableFuture + ); + + String etag = "non-matching-etag"; + ConditionalWriteOptions options = ConditionalWriteOptions.ifMatch(etag); + + AtomicReference streamRef = new AtomicReference<>(); + CompletableFuture resultFuture = asyncTransferManager.uploadObjectConditionally( + s3AsyncClient, + new UploadRequest( + "bucket", + "key", + ByteSizeUnit.MB.toBytes(1), + WritePriority.HIGH, + uploadSuccess -> {}, + false, + null, + true, + null, + options + ), + new StreamContext((partIdx, partSize, position) -> { + streamRef.set(new ZeroInputStream(partSize)); + return new InputStreamContainer(streamRef.get(), partSize, position); + }, ByteSizeUnit.MB.toBytes(1), ByteSizeUnit.MB.toBytes(1), 1), + new StatsMetricPublisher() + ); + + try { + resultFuture.get(); + fail("Expected an exception for precondition failed"); + } catch (ExecutionException | InterruptedException e) { + Throwable cause = e.getCause(); + + assertTrue("Should be S3Exception", cause instanceof S3Exception); + + S3Exception s3e = (S3Exception) cause; + + assertEquals("Should have 412 status code", 412, s3e.statusCode()); + + assertNotNull("Exception should have a message", s3e.getMessage()); + assertFalse("Exception message should not be empty", s3e.getMessage().isEmpty()); + } + + verify(s3AsyncClient, times(1)).putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class)); + + boolean closeError = false; + try { + streamRef.get().available(); + } catch (IOException e) { + closeError = e.getMessage().equals("Stream closed"); + } + assertTrue("InputStream was still open after upload", closeError); + } + + public void testConditionalOneChunkUploadCorruption() { + CompletableFuture putObjectResponseCompletableFuture = new CompletableFuture<>(); + putObjectResponseCompletableFuture.completeExceptionally( + S3Exception.builder() + .statusCode(HttpStatusCode.BAD_REQUEST) + .awsErrorDetails(AwsErrorDetails.builder().errorCode("BadDigest").build()) + .build() + ); + when(s3AsyncClient.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))).thenReturn( + putObjectResponseCompletableFuture + ); + + CompletableFuture deleteObjectResponseCompletableFuture = new CompletableFuture<>(); + deleteObjectResponseCompletableFuture.complete(DeleteObjectResponse.builder().build()); + when(s3AsyncClient.deleteObject(any(DeleteObjectRequest.class))).thenReturn(deleteObjectResponseCompletableFuture); + ConditionalWriteOptions options = ConditionalWriteOptions.ifMatch("test-etag"); + AtomicReference streamRef = new AtomicReference<>(); + CompletableFuture resultFuture = asyncTransferManager.uploadObjectConditionally( + s3AsyncClient, + new UploadRequest( + "bucket", + "key", + ByteSizeUnit.MB.toBytes(1), + WritePriority.HIGH, + uploadSuccess -> {}, + false, + null, + true, + null, + options + ), + new StreamContext((partIdx, partSize, position) -> { + streamRef.set(new ZeroInputStream(partSize)); + return new InputStreamContainer(streamRef.get(), partSize, position); + }, ByteSizeUnit.MB.toBytes(1), ByteSizeUnit.MB.toBytes(1), 1), + new StatsMetricPublisher() + ); + + try { + resultFuture.get(); + fail("Expected a corruption exception"); + } catch (ExecutionException | InterruptedException e) { + Throwable throwable = ExceptionsHelper.unwrap(e, CorruptFileException.class); + assertNotNull("Exception should be a CorruptFileException", throwable); + assertTrue("Exception should be a CorruptFileException", throwable instanceof CorruptFileException); + } + + verify(s3AsyncClient, times(1)).putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class)); + verify(s3AsyncClient, times(1)).deleteObject(any(DeleteObjectRequest.class)); + + boolean closeError = false; + try { + streamRef.get().available(); + } catch (IOException e) { + closeError = e.getMessage().equals("Stream closed"); + } + assertTrue("InputStream was still open after upload", closeError); + } + + public void testConditionalMultipartUploadPreconditionFailed() { + CompletableFuture createMultipartUploadRequestCompletableFuture = new CompletableFuture<>(); + createMultipartUploadRequestCompletableFuture.complete(CreateMultipartUploadResponse.builder().uploadId("uploadId").build()); + when(s3AsyncClient.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( + createMultipartUploadRequestCompletableFuture + ); + + CompletableFuture uploadPartResponseCompletableFuture = new CompletableFuture<>(); + uploadPartResponseCompletableFuture.complete(UploadPartResponse.builder().checksumCRC32("pzjqHA==").build()); + when(s3AsyncClient.uploadPart(any(UploadPartRequest.class), any(AsyncRequestBody.class))).thenReturn( + uploadPartResponseCompletableFuture + ); + + CompletableFuture completeMultipartUploadResponseCompletableFuture = new CompletableFuture<>(); + completeMultipartUploadResponseCompletableFuture.completeExceptionally( + S3Exception.builder().statusCode(412).message("Precondition Failed").build() + ); + when(s3AsyncClient.completeMultipartUpload(any(CompleteMultipartUploadRequest.class))).thenReturn( + completeMultipartUploadResponseCompletableFuture + ); + + CompletableFuture abortMultipartUploadResponseCompletableFuture = new CompletableFuture<>(); + abortMultipartUploadResponseCompletableFuture.complete(AbortMultipartUploadResponse.builder().build()); + when(s3AsyncClient.abortMultipartUpload(any(AbortMultipartUploadRequest.class))).thenReturn( + abortMultipartUploadResponseCompletableFuture + ); + + ConditionalWriteOptions options = ConditionalWriteOptions.ifMatch("non-matching-etag"); + + List streams = new ArrayList<>(); + CompletableFuture resultFuture = asyncTransferManager.uploadObjectConditionally( + s3AsyncClient, + new UploadRequest( + "bucket", + "key", + ByteSizeUnit.MB.toBytes(5), + WritePriority.HIGH, + uploadSuccess -> {}, + true, + 3376132981L, + true, + null, + options + ), + new StreamContext((partIdx, partSize, position) -> { + InputStream stream = new ZeroInputStream(partSize); + streams.add(stream); + return new InputStreamContainer(stream, partSize, position); + }, ByteSizeUnit.MB.toBytes(1), ByteSizeUnit.MB.toBytes(1), 5), + new StatsMetricPublisher() + ); + + try { + resultFuture.get(); + fail("Expected an exception for precondition failed"); + } catch (ExecutionException | InterruptedException e) { + Throwable cause = e.getCause(); + assertTrue("Should be S3Exception", cause instanceof S3Exception); + S3Exception s3e = (S3Exception) cause; + assertEquals("Should have 412 status code", 412, s3e.statusCode()); + assertTrue("Message should indicate condition failure", s3e.getMessage().contains("Conditional write failed")); + } + + verify(s3AsyncClient, times(1)).createMultipartUpload(any(CreateMultipartUploadRequest.class)); + verify(s3AsyncClient, times(5)).uploadPart(any(UploadPartRequest.class), any(AsyncRequestBody.class)); + verify(s3AsyncClient, times(1)).completeMultipartUpload(any(CompleteMultipartUploadRequest.class)); + verify(s3AsyncClient, times(1)).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); + + for (InputStream stream : streams) { + boolean closeError = false; + try { + stream.available(); + } catch (IOException e) { + closeError = e.getMessage().equals("Stream closed"); + } + assertTrue("InputStream was still open after upload", closeError); + } + } + + public void testConditionalMultipartUploadCorruption() { + CompletableFuture createMultipartUploadRequestCompletableFuture = new CompletableFuture<>(); + createMultipartUploadRequestCompletableFuture.complete(CreateMultipartUploadResponse.builder().uploadId("uploadId").build()); + when(s3AsyncClient.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn( + createMultipartUploadRequestCompletableFuture + ); + + CompletableFuture uploadPartResponseCompletableFuture = new CompletableFuture<>(); + uploadPartResponseCompletableFuture.complete(UploadPartResponse.builder().checksumCRC32("pzjqHA==").build()); + when(s3AsyncClient.uploadPart(any(UploadPartRequest.class), any(AsyncRequestBody.class))).thenReturn( + uploadPartResponseCompletableFuture + ); + + CompletableFuture abortMultipartUploadResponseCompletableFuture = new CompletableFuture<>(); + abortMultipartUploadResponseCompletableFuture.complete(AbortMultipartUploadResponse.builder().build()); + when(s3AsyncClient.abortMultipartUpload(any(AbortMultipartUploadRequest.class))).thenReturn( + abortMultipartUploadResponseCompletableFuture + ); + + ConditionalWriteOptions options = ConditionalWriteOptions.ifMatch("test-etag"); + + List streams = new ArrayList<>(); + CompletableFuture resultFuture = asyncTransferManager.uploadObjectConditionally( + s3AsyncClient, + new UploadRequest( + "bucket", + "key", + ByteSizeUnit.MB.toBytes(5), + WritePriority.HIGH, + uploadSuccess -> {}, + true, + 0L, + true, + null, + options + ), + new StreamContext((partIdx, partSize, position) -> { + InputStream stream = new ZeroInputStream(partSize); + streams.add(stream); + return new InputStreamContainer(stream, partSize, position); + }, ByteSizeUnit.MB.toBytes(1), ByteSizeUnit.MB.toBytes(1), 5), + new StatsMetricPublisher() + ); + + try { + resultFuture.get(); + fail("Expected a corruption exception"); + } catch (ExecutionException | InterruptedException e) { + Throwable throwable = ExceptionsHelper.unwrap(e, CorruptFileException.class); + assertNotNull("Exception should be a CorruptFileException", throwable); + assertTrue("Exception should be a CorruptFileException", throwable instanceof CorruptFileException); + } + + verify(s3AsyncClient, times(1)).createMultipartUpload(any(CreateMultipartUploadRequest.class)); + verify(s3AsyncClient, times(5)).uploadPart(any(UploadPartRequest.class), any(AsyncRequestBody.class)); + verify(s3AsyncClient, times(0)).completeMultipartUpload(any(CompleteMultipartUploadRequest.class)); + verify(s3AsyncClient, times(1)).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); + + for (InputStream stream : streams) { + boolean closeError = false; + try { + stream.available(); + } catch (IOException e) { + closeError = e.getMessage().equals("Stream closed"); + } + assertTrue("InputStream was still open after upload", closeError); + } + } } From f0857dc5918079d659a574f406dc72ea4b912576 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Wed, 28 May 2025 12:58:20 +0530 Subject: [PATCH 08/20] Test Suite for ASync Conditional upload process --- .../s3/S3BlobContainerMockClientTests.java | 216 ++++++++++++++++++ .../s3/S3BlobContainerRetriesTests.java | 65 ++++++ .../s3/S3BlobStoreContainerTests.java | 167 +++++++++----- 3 files changed, 394 insertions(+), 54 deletions(-) diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java index 9b413ac81d766..619a2f6548959 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java @@ -29,10 +29,13 @@ import software.amazon.awssdk.services.s3.model.UploadPartResponse; import org.apache.lucene.store.IndexInput; +import org.opensearch.OpenSearchException; import org.opensearch.cluster.metadata.RepositoryMetadata; import org.opensearch.common.CheckedConsumer; import org.opensearch.common.StreamContext; import org.opensearch.common.blobstore.BlobPath; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteOptions; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteResponse; import org.opensearch.common.blobstore.stream.write.StreamContextSupplier; import org.opensearch.common.blobstore.stream.write.WriteContext; import org.opensearch.common.blobstore.stream.write.WritePriority; @@ -75,6 +78,7 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.stream.IntStream; +import org.mockito.ArgumentMatchers; import org.mockito.Mockito; import org.mockito.invocation.InvocationOnMock; @@ -84,6 +88,7 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.any; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doCallRealMethod; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; @@ -756,4 +761,215 @@ private void testLargeFilesRedirectedToSlowSyncClient(boolean expectException, W } }); } + + private void testLargeFilesRedirectedToSlowSyncClientConditional( + boolean expectException, + WritePriority writePriority, + int conditionalResponseCode + ) throws IOException, InterruptedException { + + ByteSizeValue capacity = new ByteSizeValue(1, ByteSizeUnit.GB); + int numberOfParts = 20; + final ByteSizeValue partSize = new ByteSizeValue(capacity.getBytes() / numberOfParts + 1, ByteSizeUnit.BYTES); + + GenericStatsMetricPublisher genericStatsMetricPublisher = new GenericStatsMetricPublisher(10000L, 10, 10000L, 10); + SizeBasedBlockingQ sizeBasedBlockingQ = new SizeBasedBlockingQ( + capacity, + transferQueueConsumerService, + 10, + genericStatsMetricPublisher, + SizeBasedBlockingQ.QueueEventType.NORMAL + ); + + final long lastPartSize = new ByteSizeValue(200, ByteSizeUnit.MB).getBytes(); + final long blobSize = ((numberOfParts - 1) * partSize.getBytes()) + lastPartSize; + CountDownLatch countDownLatch = new CountDownLatch(1); + AtomicReference responseRef = new AtomicReference<>(); + AtomicReference exceptionRef = new AtomicReference<>(); + + ActionListener completionListener = ActionListener.wrap(resp -> { + responseRef.set(resp); + countDownLatch.countDown(); + }, ex -> { + exceptionRef.set(ex); + countDownLatch.countDown(); + }); + + final String bucketName = randomAlphaOfLengthBetween(1, 10); + + final BlobPath blobPath = new BlobPath(); + if (randomBoolean()) { + IntStream.of(randomIntBetween(1, 5)).forEach(value -> blobPath.add("path_" + value)); + } + + final long bufferSize = ByteSizeUnit.MB.toBytes(randomIntBetween(5, 1024)); + + final S3BlobStore blobStore = mock(S3BlobStore.class); + when(blobStore.bucket()).thenReturn(bucketName); + when(blobStore.getStatsMetricPublisher()).thenReturn(new StatsMetricPublisher()); + when(blobStore.bufferSizeInBytes()).thenReturn(bufferSize); + + when(blobStore.getLowPrioritySizeBasedBlockingQ()).thenReturn(sizeBasedBlockingQ); + when(blobStore.getNormalPrioritySizeBasedBlockingQ()).thenReturn(sizeBasedBlockingQ); + + final StorageClass storageClass = randomFrom(StorageClass.values()); + when(blobStore.getStorageClass()).thenReturn(storageClass); + when(blobStore.isRedirectLargeUploads()).thenReturn(true); + boolean uploadRetryEnabled = randomBoolean(); + when(blobStore.isUploadRetryEnabled()).thenReturn(uploadRetryEnabled); + + final ObjectCannedACL cannedAccessControlList = randomBoolean() ? randomFrom(ObjectCannedACL.values()) : null; + if (cannedAccessControlList != null) { + when(blobStore.getCannedACL()).thenReturn(cannedAccessControlList); + } + + final S3Client client = mock(S3Client.class); + final AmazonS3Reference clientReference = Mockito.spy(new AmazonS3Reference(client)); + doNothing().when(clientReference).close(); + when(blobStore.clientReference()).thenReturn(clientReference); + + final String uploadId = randomAlphaOfLength(10); + final CreateMultipartUploadResponse createMultipartUploadResponse = CreateMultipartUploadResponse.builder() + .uploadId(uploadId) + .build(); + when(client.createMultipartUpload(any(CreateMultipartUploadRequest.class))).thenReturn(createMultipartUploadResponse); + + if (expectException) { + when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenThrow( + SdkException.create("Expected upload part request to fail", new RuntimeException()) + ); + } else { + when(client.uploadPart(any(UploadPartRequest.class), any(RequestBody.class))).thenReturn( + UploadPartResponse.builder().eTag("part-etag-" + randomAlphaOfLength(5)).build() + ); + } + + if (conditionalResponseCode == 412) { + when(client.completeMultipartUpload(any(CompleteMultipartUploadRequest.class))).thenThrow( + software.amazon.awssdk.services.s3.model.S3Exception.builder().statusCode(412).message("Precondition Failed").build() + ); + } else if (conditionalResponseCode == 409) { + when(client.completeMultipartUpload(any(CompleteMultipartUploadRequest.class))).thenThrow( + software.amazon.awssdk.services.s3.model.S3Exception.builder().statusCode(409).message("Resource Already Exists").build() + ); + } else { + String eTag = "\"multipart-etag-" + randomAlphaOfLength(5) + "\""; + when(client.completeMultipartUpload(any(CompleteMultipartUploadRequest.class))).thenReturn( + CompleteMultipartUploadResponse.builder().eTag(eTag).build() + ); + } + + when(client.abortMultipartUpload(any(AbortMultipartUploadRequest.class))).thenReturn( + AbortMultipartUploadResponse.builder().build() + ); + + List openInputStreams = new ArrayList<>(); + final S3BlobContainer s3BlobContainer = Mockito.spy(new S3BlobContainer(blobPath, blobStore)); + + doCallRealMethod().when(s3BlobContainer) + .executeMultipartUploadConditionally( + any(S3BlobStore.class), + anyString(), + any(InputStream.class), + anyLong(), + any(), + any(ConditionalWriteOptions.class), + ArgumentMatchers.>any() + ); + + StreamContextSupplier streamContextSupplier = partSize1 -> new StreamContext((partNo, size, position) -> { + InputStream inputStream = new OffsetRangeIndexInputStream(new ZeroIndexInput("desc", blobSize), size, position); + openInputStreams.add(inputStream); + return new InputStreamContainer(inputStream, size, position); + }, partSize1, calculateLastPartSize(blobSize, partSize1), calculateNumberOfParts(blobSize, partSize1)); + + WriteContext writeContext = new WriteContext.Builder().fileName("write_large_blob_conditional") + .streamContextSupplier(streamContextSupplier) + .fileSize(blobSize) + .failIfAlreadyExists(false) + .writePriority(writePriority) + .uploadFinalizer(success -> { + Assert.assertTrue(success); + }) + .doRemoteDataIntegrityCheck(false) + .metadata(new HashMap<>()) + .build(); + + ConditionalWriteOptions conditionalOptions; + if (conditionalResponseCode == 412) { + conditionalOptions = ConditionalWriteOptions.ifMatch("invalid-etag"); + } else if (conditionalResponseCode == 409) { + conditionalOptions = ConditionalWriteOptions.ifNotExists(); + } else { + conditionalOptions = ConditionalWriteOptions.ifMatch("valid-etag"); + } + + s3BlobContainer.asyncBlobUploadConditionally(writeContext, conditionalOptions, completionListener); + + boolean awaitSuccess = countDownLatch.await(5000, TimeUnit.SECONDS); + assertTrue(awaitSuccess); + + if (expectException || conditionalResponseCode != 0) { + assertNotNull("Should have received an exception", exceptionRef.get()); + + if (conditionalResponseCode == 412 && !expectException) { + assertTrue( + "Should have conditional error", + exceptionRef.get() instanceof OpenSearchException + || exceptionRef.get().getMessage().contains("Precondition Failed") + || exceptionRef.get().getMessage().contains("ETag mismatch") + ); + } else if (conditionalResponseCode == 409 && !expectException) { + assertTrue( + "Should have conflict error", + exceptionRef.get().getMessage().contains("Resource Already Exists") + || exceptionRef.get().getMessage().contains("already exists") + ); + } + } else { + assertNull("Should not have received an exception", exceptionRef.get()); + assertNotNull("Should have received a response", responseRef.get()); + assertNotNull("Response should have version identifier", responseRef.get().getVersionIdentifier()); + } + + verify(s3BlobContainer, times(1)).executeMultipartUploadConditionally( + any(S3BlobStore.class), + anyString(), + any(InputStream.class), + anyLong(), + anyMap(), + any(ConditionalWriteOptions.class), + ArgumentMatchers.>any() + ); + + boolean shouldAbort = expectException || (conditionalResponseCode != 0 && !expectException); + verify(client, times(shouldAbort ? 1 : 0)).abortMultipartUpload(any(AbortMultipartUploadRequest.class)); + + openInputStreams.forEach(inputStream -> { + try { + inputStream.close(); + } catch (IOException ex) {} + }); + } + + public void testFailureWhenLargeFileRedirectedConditional() throws IOException, InterruptedException { + testLargeFilesRedirectedToSlowSyncClientConditional(true, WritePriority.LOW, 0); + testLargeFilesRedirectedToSlowSyncClientConditional(true, WritePriority.NORMAL, 0); + } + + public void testLargeFileRedirectedConditional() throws IOException, InterruptedException { + testLargeFilesRedirectedToSlowSyncClientConditional(false, WritePriority.LOW, 0); + testLargeFilesRedirectedToSlowSyncClientConditional(false, WritePriority.NORMAL, 0); + } + + public void testLargeFileRedirectedConditionalPreconditionFailed() throws IOException, InterruptedException { + testLargeFilesRedirectedToSlowSyncClientConditional(false, WritePriority.LOW, 412); + testLargeFilesRedirectedToSlowSyncClientConditional(false, WritePriority.NORMAL, 412); + } + + public void testLargeFileRedirectedConditionalConflict() throws IOException, InterruptedException { + testLargeFilesRedirectedToSlowSyncClientConditional(false, WritePriority.LOW, 409); + testLargeFilesRedirectedToSlowSyncClientConditional(false, WritePriority.NORMAL, 409); + } + } diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerRetriesTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerRetriesTests.java index 96ef28d24c14f..849c3acd21599 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerRetriesTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerRetriesTests.java @@ -33,6 +33,7 @@ import software.amazon.awssdk.core.exception.SdkClientException; import software.amazon.awssdk.core.io.SdkDigestInputStream; +import software.amazon.awssdk.services.s3.model.S3Exception; import software.amazon.awssdk.utils.internal.Base16; import org.apache.http.HttpStatus; @@ -43,6 +44,8 @@ import org.opensearch.common.blobstore.AsyncMultiStreamBlobContainer; import org.opensearch.common.blobstore.BlobContainer; import org.opensearch.common.blobstore.BlobPath; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteOptions; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteResponse; import org.opensearch.common.blobstore.stream.write.StreamContextSupplier; import org.opensearch.common.blobstore.stream.write.WriteContext; import org.opensearch.common.blobstore.stream.write.WritePriority; @@ -420,6 +423,68 @@ public void testWriteBlobByStreamsWithRetries() throws Exception { }); } + public void testWriteBlobByStreamsConditionalFailure() throws Exception { + byte[] bytes = randomBlobContent(); + + httpServer.createContext("/bucket/write_blob_conditionally_failure", exchange -> { + if ("PUT".equals(exchange.getRequestMethod()) && exchange.getRequestURI().getQuery() == null) { + Streams.readFully(exchange.getRequestBody()); + exchange.sendResponseHeaders(HttpStatus.SC_PRECONDITION_FAILED, -1); + exchange.close(); + } + }); + + AsyncMultiStreamBlobContainer blobContainer = createBlobContainer(0, null, true, null); + List streams = new ArrayList<>(); + CountDownLatch latch = new CountDownLatch(1); + AtomicReference exRef = new AtomicReference<>(); + + ActionListener listener = ActionListener.wrap(resp -> { + fail("Expected failure"); + latch.countDown(); + }, ex -> { + exRef.set(ex); + latch.countDown(); + }); + + int partSize = 5 * 1024 * 1024; + StreamContextSupplier supplier = ps -> new StreamContext((part, size, pos) -> { + InputStream in = new OffsetRangeIndexInputStream(new ByteArrayIndexInput("desc", bytes), size, pos); + streams.add(in); + return new InputStreamContainer(in, size, pos); + }, partSize, calculateLastPartSize(bytes.length, partSize), calculateNumberOfParts(bytes.length, partSize)); + + ConditionalWriteOptions options = ConditionalWriteOptions.ifMatch("non-matching-etag"); + WriteContext context = new WriteContext.Builder().fileName("write_blob_conditionally_failure") + .streamContextSupplier(supplier) + .fileSize(bytes.length) + .failIfAlreadyExists(false) + .writePriority(WritePriority.NORMAL) + .uploadFinalizer(Assert::assertTrue) + .doRemoteDataIntegrityCheck(false) + .build(); + + blobContainer.asyncBlobUploadConditionally(context, options, listener); + assertTrue(latch.await(5, TimeUnit.SECONDS)); + + Exception actual = exRef.get(); + assertNotNull(actual); + String msg = actual.getMessage(); + assertTrue( + msg.contains("Precondition Failed") + || msg.contains("412") + || (actual instanceof S3Exception && ((S3Exception) actual).statusCode() == 412) + ); + + for (InputStream in : streams) { + try { + in.close(); + } catch (IOException e) { + fail("Stream close failure"); + } + } + } + private long calculateLastPartSize(long totalSize, long partSize) { return totalSize % partSize == 0 ? partSize : totalSize % partSize; } diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java index 4e9be32972e36..d360f95a4e8d8 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java @@ -79,6 +79,8 @@ import org.opensearch.common.blobstore.BlobMetadata; import org.opensearch.common.blobstore.BlobPath; import org.opensearch.common.blobstore.BlobStoreException; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteOptions; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteResponse; import org.opensearch.common.blobstore.DeleteResult; import org.opensearch.common.blobstore.stream.read.ReadContext; import org.opensearch.common.collect.Tuple; @@ -944,7 +946,7 @@ public void testNumberOfMultiparts() { assertNumberOfMultiparts(factor + 1, remaining, (size * factor) + remaining, size); } - public void testExecuteMultipartUploadIfEtagMatchesSuccess() throws IOException { + public void testExecuteMultipartUploadConditionallyWithEtagMatchSuccess() throws IOException { final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String inputETag = randomAlphaOfLengthBetween(8, 32); @@ -1021,10 +1023,19 @@ public void testExecuteMultipartUploadIfEtagMatchesSuccess() throws IOException ); @SuppressWarnings("unchecked") - ActionListener etagListener = mock(ActionListener.class); + ActionListener responseListener = mock(ActionListener.class); final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); - blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, metadata, inputETag, etagListener); + + blobContainer.executeMultipartUploadConditionally( + blobStore, + blobName, + inputStream, + blobSize, + metadata, + ConditionalWriteOptions.ifMatch(inputETag), + responseListener + ); final CreateMultipartUploadRequest createRequest = createRequestCaptor.getValue(); assertEquals(bucketName, createRequest.bucket()); @@ -1067,13 +1078,15 @@ public void testExecuteMultipartUploadIfEtagMatchesSuccess() throws IOException assertEquals(inputETag, completeRequest.ifMatch()); - verify(etagListener).onResponse(finalETag); - verify(etagListener, never()).onFailure(any()); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(ConditionalWriteResponse.class); + verify(responseListener).onResponse(responseCaptor.capture()); + assertEquals(finalETag, responseCaptor.getValue().getVersionIdentifier()); + verify(responseListener, never()).onFailure(any()); verify(clientReference).close(); } - public void testExecuteMultipartUploadIfEtagMatchesWithMetadataAndSSE() throws IOException { + public void testExecuteMultipartUploadConditionallyWithMetadataAndSSE() throws IOException { final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String inputETag = randomAlphaOfLengthBetween(8, 32); @@ -1140,10 +1153,19 @@ public void testExecuteMultipartUploadIfEtagMatchesWithMetadataAndSSE() throws I ); @SuppressWarnings("unchecked") - ActionListener etagListener = mock(ActionListener.class); + ActionListener responseListener = mock(ActionListener.class); final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); - blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, metadata, inputETag, etagListener); + + blobContainer.executeMultipartUploadConditionally( + blobStore, + blobName, + inputStream, + blobSize, + metadata, + ConditionalWriteOptions.ifMatch(inputETag), + responseListener + ); final CreateMultipartUploadRequest createRequest = createRequestCaptor.getValue(); assertEquals(bucketName, createRequest.bucket()); @@ -1181,14 +1203,16 @@ public void testExecuteMultipartUploadIfEtagMatchesWithMetadataAndSSE() throws I assertEquals(inputETag, completeRequest.ifMatch()); - verify(etagListener).onResponse(finalETag); - verify(etagListener, never()).onFailure(any()); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(ConditionalWriteResponse.class); + verify(responseListener).onResponse(responseCaptor.capture()); + assertEquals(finalETag, responseCaptor.getValue().getVersionIdentifier()); - verify(clientReference).close(); + verify(responseListener, never()).onFailure(any()); + verify(clientReference).close(); } - public void testExecuteMultipartUploadIfEtagMatchesContentIntegrity() throws IOException { + public void testExecuteMultipartUploadConditionallyContentIntegrity() throws IOException { final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String inputETag = randomAlphaOfLengthBetween(8, 32); @@ -1251,10 +1275,19 @@ public void testExecuteMultipartUploadIfEtagMatchesContentIntegrity() throws IOE ); @SuppressWarnings("unchecked") - ActionListener etagListener = mock(ActionListener.class); + ActionListener responseListener = mock(ActionListener.class); final ByteArrayInputStream inputStream = new ByteArrayInputStream(blobContent); - blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, null, inputETag, etagListener); + + blobContainer.executeMultipartUploadConditionally( + blobStore, + blobName, + inputStream, + blobSize, + null, + ConditionalWriteOptions.ifMatch(inputETag), + responseListener + ); final CreateMultipartUploadRequest createRequest = createRequestCaptor.getValue(); assertEquals(bucketName, createRequest.bucket()); @@ -1277,13 +1310,16 @@ public void testExecuteMultipartUploadIfEtagMatchesContentIntegrity() throws IOE assertArrayEquals("Uploaded content should match original content", blobContent, reassembledContent); - verify(etagListener).onResponse(finalETag); - verify(etagListener, never()).onFailure(any()); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(ConditionalWriteResponse.class); + verify(responseListener).onResponse(responseCaptor.capture()); + assertEquals(finalETag, responseCaptor.getValue().getVersionIdentifier()); + + verify(responseListener, never()).onFailure(any()); verify(clientReference).close(); } - public void testExecuteMultipartUploadIfEtagMatchesSizeValidation() { + public void testExecuteMultipartUploadConditionallySizeValidation() { final S3BlobStore blobStore = mock(S3BlobStore.class); final S3BlobContainer blobContainer = new S3BlobContainer(mock(BlobPath.class), blobStore); final String blobName = randomAlphaOfLengthBetween(1, 10); @@ -1291,20 +1327,20 @@ public void testExecuteMultipartUploadIfEtagMatchesSizeValidation() { final String finalETag = randomAlphaOfLengthBetween(8, 32); @SuppressWarnings("unchecked") - ActionListener invalidSizeListener = mock(ActionListener.class); + ActionListener invalidSizeListener = mock(ActionListener.class); { final long tooSmallSize = ByteSizeUnit.MB.toBytes(5) - 1024; final IllegalArgumentException tooSmallException = expectThrows( IllegalArgumentException.class, - () -> blobContainer.executeMultipartUploadIfEtagMatches( + () -> blobContainer.executeMultipartUploadConditionally( blobStore, blobName, new ByteArrayInputStream(new byte[0]), tooSmallSize, null, - inputETag, + ConditionalWriteOptions.ifMatch(inputETag), invalidSizeListener ) ); @@ -1319,13 +1355,13 @@ public void testExecuteMultipartUploadIfEtagMatchesSizeValidation() { final IllegalArgumentException tooLargeException = expectThrows( IllegalArgumentException.class, - () -> blobContainer.executeMultipartUploadIfEtagMatches( + () -> blobContainer.executeMultipartUploadConditionally( blobStore, blobName, new ByteArrayInputStream(new byte[0]), tooLargeSize, null, - inputETag, + ConditionalWriteOptions.ifMatch(inputETag), invalidSizeListener ) ); @@ -1365,7 +1401,7 @@ public void testExecuteMultipartUploadIfEtagMatchesSizeValidation() { { final long exactMinimumSize = ByteSizeUnit.MB.toBytes(5); @SuppressWarnings("unchecked") - ActionListener validSizeListener = mock(ActionListener.class); + ActionListener validSizeListener = mock(ActionListener.class); InputStream zeroStream = new InputStream() { long remaining = exactMinimumSize; @@ -1393,17 +1429,19 @@ public int read(byte[] b, int off, int len) { }; try { - blobContainer.executeMultipartUploadIfEtagMatches( + blobContainer.executeMultipartUploadConditionally( blobStore, blobName, zeroStream, exactMinimumSize, null, - inputETag, + ConditionalWriteOptions.ifMatch(inputETag), validSizeListener ); - verify(validSizeListener).onResponse(finalETag); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(ConditionalWriteResponse.class); + verify(validSizeListener).onResponse(responseCaptor.capture()); + assertEquals(finalETag, responseCaptor.getValue().getVersionIdentifier()); verify(validSizeListener, never()).onFailure(any()); verify(clientReference).close(); @@ -1423,7 +1461,7 @@ public int read(byte[] b, int off, int len) { { final long testSize = ByteSizeUnit.MB.toBytes(10); @SuppressWarnings("unchecked") - ActionListener validSizeListener = mock(ActionListener.class); + ActionListener validSizeListener = mock(ActionListener.class); InputStream zeroStream = new InputStream() { long remaining = testSize; @@ -1451,17 +1489,19 @@ public int read(byte[] b, int off, int len) { }; try { - blobContainer.executeMultipartUploadIfEtagMatches( + blobContainer.executeMultipartUploadConditionally( blobStore, blobName, zeroStream, testSize, null, - inputETag, + ConditionalWriteOptions.ifMatch(inputETag), validSizeListener ); - verify(validSizeListener).onResponse(finalETag); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(ConditionalWriteResponse.class); + verify(validSizeListener).onResponse(responseCaptor.capture()); + assertEquals(finalETag, responseCaptor.getValue().getVersionIdentifier()); verify(validSizeListener, never()).onFailure(any()); verify(clientReference).close(); @@ -1475,7 +1515,7 @@ public int read(byte[] b, int off, int len) { } } - public void testExecuteMultipartUploadIfEtagMatchesPreconditionFailed() { + public void testExecuteMultipartUploadConditionallyPreconditionFailed() { final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String eTag = randomAlphaOfLengthBetween(8, 32); @@ -1524,7 +1564,7 @@ public void testExecuteMultipartUploadIfEtagMatchesPreconditionFailed() { ); final AtomicReference capturedException = new AtomicReference<>(); - ActionListener etagListener = ActionListener.wrap( + ActionListener responseListener = ActionListener.wrap( r -> fail("Should have failed with precondition failure"), capturedException::set ); @@ -1533,7 +1573,15 @@ public void testExecuteMultipartUploadIfEtagMatchesPreconditionFailed() { IOException ioException = expectThrows( IOException.class, - () -> blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, null, eTag, etagListener) + () -> blobContainer.executeMultipartUploadConditionally( + blobStore, + blobName, + inputStream, + blobSize, + null, + ConditionalWriteOptions.ifMatch(eTag), + responseListener + ) ); assertEquals("Unable to upload object [" + blobName + "] due to ETag mismatch", ioException.getMessage()); @@ -1542,7 +1590,7 @@ public void testExecuteMultipartUploadIfEtagMatchesPreconditionFailed() { Exception exception = capturedException.get(); assertNotNull("Expected an exception to be captured", exception); assertTrue("Exception should be an OpenSearchException", exception instanceof OpenSearchException); - assertEquals("stale_primary_shard", ((OpenSearchException) exception).getMessage()); + assertEquals("stale_primary_shard", exception.getMessage()); verify(client).createMultipartUpload(any(CreateMultipartUploadRequest.class)); verify(client).completeMultipartUpload(any(CompleteMultipartUploadRequest.class)); @@ -1552,7 +1600,7 @@ public void testExecuteMultipartUploadIfEtagMatchesPreconditionFailed() { verify(abortClientReference).close(); } - public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { + public void testExecuteMultipartUploadConditionallyS3ExceptionTypes() { final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String eTag = randomAlphaOfLengthBetween(8, 32); @@ -1630,7 +1678,8 @@ public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { ); final AtomicReference capturedException = new AtomicReference<>(); - ActionListener etagListener = ActionListener.wrap( + ActionListener responseListener = ActionListener.wrap( + // Changed listener type r -> fail("Should have failed with exception"), capturedException::set ); @@ -1638,14 +1687,14 @@ public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); IOException exception = expectThrows( IOException.class, - () -> blobContainer.executeMultipartUploadIfEtagMatches( + () -> blobContainer.executeMultipartUploadConditionally( blobStore, blobName, inputStream, blobSize, metadata, - eTag, - etagListener + ConditionalWriteOptions.ifMatch(eTag), + responseListener ) ); @@ -1672,7 +1721,7 @@ public void testExecuteMultipartUploadIfEtagMatchesS3ExceptionTypes() { } } - public void testExecuteMultipartUploadIfEtagMatchesSdkException() { + public void testExecuteMultipartUploadConditionallySdkException() { final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String eTag = randomAlphaOfLengthBetween(8, 32); @@ -1730,7 +1779,7 @@ public void testExecuteMultipartUploadIfEtagMatchesSdkException() { } final AtomicReference capturedException = new AtomicReference<>(); - ActionListener etagListener = ActionListener.wrap( + ActionListener responseListener = ActionListener.wrap( r -> fail("Should have failed with SdkException"), capturedException::set ); @@ -1738,14 +1787,14 @@ public void testExecuteMultipartUploadIfEtagMatchesSdkException() { final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); IOException exception = expectThrows( IOException.class, - () -> blobContainer.executeMultipartUploadIfEtagMatches( + () -> blobContainer.executeMultipartUploadConditionally( blobStore, blobName, inputStream, blobSize, null, - eTag, - etagListener + ConditionalWriteOptions.ifMatch(eTag), + responseListener ) ); @@ -1773,7 +1822,7 @@ public void testExecuteMultipartUploadIfEtagMatchesSdkException() { } } - public void testExecuteMultipartUploadIfEtagMatchesResourceManagement() throws IOException { + public void testExecuteMultipartUploadConditionallyResourceManagement() throws IOException { final String bucketName = randomAlphaOfLengthBetween(1, 10); final String blobName = randomAlphaOfLengthBetween(1, 10); final String eTag = randomAlphaOfLengthBetween(8, 32); @@ -1863,36 +1912,46 @@ enum ResourceScenario { } @SuppressWarnings("unchecked") - ActionListener etagListener = mock(ActionListener.class); + ActionListener responseListener = mock(ActionListener.class); final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[(int) blobSize]); if (scenario == ResourceScenario.SUCCESS_PATH) { - blobContainer.executeMultipartUploadIfEtagMatches(blobStore, blobName, inputStream, blobSize, null, eTag, etagListener); + blobContainer.executeMultipartUploadConditionally( + blobStore, + blobName, + inputStream, + blobSize, + null, + ConditionalWriteOptions.ifMatch(eTag), + responseListener + ); - verify(etagListener).onResponse("final-etag"); - verify(etagListener, never()).onFailure(any()); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(ConditionalWriteResponse.class); + verify(responseListener).onResponse(responseCaptor.capture()); + assertEquals("final-etag", responseCaptor.getValue().getVersionIdentifier()); + verify(responseListener, never()).onFailure(any()); verify(blobStore, times(1)).clientReference(); verify(primaryClientReference).close(); } else { IOException exception = expectThrows( IOException.class, - () -> blobContainer.executeMultipartUploadIfEtagMatches( + () -> blobContainer.executeMultipartUploadConditionally( blobStore, blobName, inputStream, blobSize, null, - eTag, - etagListener + ConditionalWriteOptions.ifMatch(eTag), + responseListener ) ); assertEquals("Exception cause should be the original exception", stageException, exception.getCause()); - verify(etagListener).onFailure(any(Exception.class)); - verify(etagListener, never()).onResponse(any()); + verify(responseListener).onFailure(any(Exception.class)); + verify(responseListener, never()).onResponse(any()); verify(primaryClientReference).close(); From 062a76d0d17229a9af4ba2382f4bc6bcbbe1e934 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Wed, 28 May 2025 12:59:29 +0530 Subject: [PATCH 09/20] S3 Blobcontainer Implementation for Async Conditional upload --- .../repositories/s3/S3BlobContainer.java | 175 ++++++++++++++++-- 1 file changed, 157 insertions(+), 18 deletions(-) diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java index 55a170173c4ab..1592ec1e31d53 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java @@ -76,6 +76,8 @@ import org.opensearch.common.blobstore.BlobMetadata; import org.opensearch.common.blobstore.BlobPath; import org.opensearch.common.blobstore.BlobStoreException; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteOptions; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteResponse; import org.opensearch.common.blobstore.DeleteResult; import org.opensearch.common.blobstore.InputStreamWithMetadata; import org.opensearch.common.blobstore.stream.read.ReadContext; @@ -208,6 +210,117 @@ public void writeBlobWithMetadata( }); } + @Override + public void asyncBlobUploadConditionally( + WriteContext writeContext, + ConditionalWriteOptions options, + ActionListener completionListener + ) throws IOException { + UploadRequest uploadRequest = new UploadRequest( + blobStore.bucket(), + buildKey(writeContext.getFileName()), + writeContext.getFileSize(), + writeContext.getWritePriority(), + writeContext.getUploadFinalizer(), + writeContext.doRemoteDataIntegrityCheck(), + writeContext.getExpectedChecksum(), + blobStore.isUploadRetryEnabled(), + writeContext.getMetadata(), + options + ); + + try { + // If file size is greater than the queue capacity than SizeBasedBlockingQ will always reject the upload. + // Therefore, redirecting it to slow client. + if ((uploadRequest.getWritePriority() == WritePriority.LOW + && blobStore.getLowPrioritySizeBasedBlockingQ().isMaxCapacityBelowContentLength(uploadRequest.getContentLength()) == false) + || (uploadRequest.getWritePriority() != WritePriority.HIGH + && uploadRequest.getWritePriority() != WritePriority.URGENT + && blobStore.getNormalPrioritySizeBasedBlockingQ() + .isMaxCapacityBelowContentLength(uploadRequest.getContentLength()) == false)) { + + StreamContext streamContext = SocketAccess.doPrivileged( + () -> writeContext.getStreamProvider(uploadRequest.getContentLength()) + ); + InputStreamContainer inputStream = streamContext.provideStream(0); + + try { + executeMultipartUploadConditionally( + blobStore, + buildKey(writeContext.getFileName()), + inputStream.getInputStream(), + uploadRequest.getContentLength(), + uploadRequest.getMetadata(), + options, + completionListener + ); + } catch (Exception ex) { + logger.error( + "Failed to conditionally upload large file {} of size {}", + uploadRequest.getKey(), + uploadRequest.getContentLength(), + ex + ); + completionListener.onFailure(ex); + } + return; + } + + long partSize = blobStore.getAsyncTransferManager() + .calculateOptimalPartSize(writeContext.getFileSize(), writeContext.getWritePriority(), blobStore.isUploadRetryEnabled()); + + StreamContext streamContext = SocketAccess.doPrivileged(() -> writeContext.getStreamProvider(partSize)); + + try (AmazonAsyncS3Reference amazonS3Reference = SocketAccess.doPrivileged(blobStore::asyncClientReference)) { + S3AsyncClient s3AsyncClient; + if (writeContext.getWritePriority() == WritePriority.URGENT) { + s3AsyncClient = amazonS3Reference.get().urgentClient(); + } else if (writeContext.getWritePriority() == WritePriority.HIGH) { + s3AsyncClient = amazonS3Reference.get().priorityClient(); + } else { + s3AsyncClient = amazonS3Reference.get().client(); + } + + if (writeContext.getWritePriority() == WritePriority.URGENT + || writeContext.getWritePriority() == WritePriority.HIGH + || blobStore.isPermitBackedTransferEnabled() == false) { + createFileCompletableFutureConditionally(s3AsyncClient, uploadRequest, streamContext, completionListener); + } else if (writeContext.getWritePriority() == WritePriority.LOW) { + blobStore.getLowPrioritySizeBasedBlockingQ() + .produce( + new SizeBasedBlockingQ.Item( + writeContext.getFileSize(), + () -> createFileCompletableFutureConditionally( + s3AsyncClient, + uploadRequest, + streamContext, + completionListener + ) + ) + ); + } else if (writeContext.getWritePriority() == WritePriority.NORMAL) { + blobStore.getNormalPrioritySizeBasedBlockingQ() + .produce( + new SizeBasedBlockingQ.Item( + writeContext.getFileSize(), + () -> createFileCompletableFutureConditionally( + s3AsyncClient, + uploadRequest, + streamContext, + completionListener + ) + ) + ); + } else { + throw new IllegalStateException("Cannot perform upload for other priority types."); + } + } + } catch (Exception e) { + logger.info("Exception during conditional blob upload for file {}", writeContext.getFileName(), e); + throw new IOException(e); + } + } + @Override public void asyncBlobUpload(WriteContext writeContext, ActionListener completionListener) throws IOException { UploadRequest uploadRequest = new UploadRequest( @@ -219,7 +332,8 @@ public void asyncBlobUpload(WriteContext writeContext, ActionListener comp writeContext.doRemoteDataIntegrityCheck(), writeContext.getExpectedChecksum(), blobStore.isUploadRetryEnabled(), - writeContext.getMetadata() + writeContext.getMetadata(), + null ); try { // If file size is greater than the queue capacity than SizeBasedBlockingQ will always reject the upload. @@ -318,6 +432,25 @@ private CompletableFuture createFileCompletableFuture( }); } + private CompletableFuture createFileCompletableFutureConditionally( + S3AsyncClient s3AsyncClient, + UploadRequest uploadRequest, + StreamContext streamContext, + ActionListener completionListener + ) { + CompletableFuture completableFuture = blobStore.getAsyncTransferManager() + .uploadObjectConditionally(s3AsyncClient, uploadRequest, streamContext, blobStore.getStatsMetricPublisher()); + + return completableFuture.whenComplete((response, throwable) -> { + if (throwable == null) { + completionListener.onResponse(response); + } else { + Exception ex = throwable instanceof Error ? new Exception(throwable) : (Exception) throwable; + completionListener.onFailure(ex); + } + }); + } + @ExperimentalApi @Override public void readBlobAsync(String blobName, ActionListener listener) { @@ -514,14 +647,14 @@ private String buildKey(String blobName) { return keyPath + blobName; } - public void executeMultipartUploadIfEtagMatches( + public void executeMultipartUploadConditionally( final S3BlobStore blobStore, final String blobName, final InputStream input, final long blobSize, final Map metadata, - final String eTag, - final ActionListener etagListener + final ConditionalWriteOptions options, + final ActionListener listener ) throws IOException { ensureMultiPartUploadSize(blobSize); @@ -534,7 +667,7 @@ public void executeMultipartUploadIfEtagMatches( final int nbParts = multiparts.v1().intValue(); final long lastPartSize = multiparts.v2(); assert blobSize == (((nbParts - 1) * partSize) + lastPartSize) : "blobSize does not match multipart sizes"; - // test + CreateMultipartUploadRequest.Builder createRequestBuilder = CreateMultipartUploadRequest.builder() .bucket(blobStore.bucket()) .key(blobName) @@ -564,7 +697,7 @@ public void executeMultipartUploadIfEtagMatches( ); if (Strings.isEmpty(uploadId.get())) { IOException exception = new IOException("Failed to initialize multipart upload for " + blobName); - etagListener.onFailure(exception); + listener.onFailure(exception); throw exception; } @@ -594,7 +727,7 @@ public void executeMultipartUploadIfEtagMatches( IOException exception = new IOException( String.format(Locale.ROOT, "S3 part upload for [%s] part [%d] returned null ETag", blobName, i) ); - etagListener.onFailure(exception); + listener.onFailure(exception); throw exception; } @@ -605,18 +738,24 @@ public void executeMultipartUploadIfEtagMatches( IOException exception = new IOException( String.format(Locale.ROOT, "Multipart upload for [%s] sent %d bytes; expected %d bytes", blobName, bytesCount, blobSize) ); - etagListener.onFailure(exception); + listener.onFailure(exception); throw exception; } - CompleteMultipartUploadRequest completeRequest = CompleteMultipartUploadRequest.builder() + CompleteMultipartUploadRequest.Builder completeRequestBuilder = CompleteMultipartUploadRequest.builder() .bucket(bucketName) .key(blobName) .uploadId(uploadId.get()) .multipartUpload(CompletedMultipartUpload.builder().parts(parts).build()) - .ifMatch(eTag) - .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector)) - .build(); + .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector)); + + if (options.isIfMatch()) { + completeRequestBuilder.ifMatch(options.getVersionIdentifier()); + } else if (options.isIfNotExists()) { + completeRequestBuilder.ifNoneMatch("*"); + } + + CompleteMultipartUploadRequest completeRequest = completeRequestBuilder.build(); CompleteMultipartUploadResponse completeResponse = SocketAccess.doPrivileged( () -> clientReference.get().completeMultipartUpload(completeRequest) @@ -624,37 +763,37 @@ public void executeMultipartUploadIfEtagMatches( if (completeResponse.eTag() != null) { success = true; - etagListener.onResponse(completeResponse.eTag()); + listener.onResponse(ConditionalWriteResponse.success(completeResponse.eTag())); } else { IOException exception = new IOException( "S3 multipart upload for [" + blobName + "] returned null ETag, violating data integrity expectations" ); - etagListener.onFailure(exception); + listener.onFailure(exception); throw exception; } } catch (S3Exception e) { if (e.statusCode() == HTTP_STATUS_PRECONDITION_FAILED) { - etagListener.onFailure(new OpenSearchException("stale_primary_shard", e, "Precondition Failed : Etag Mismatch", blobName)); + listener.onFailure(new OpenSearchException("stale_primary_shard", e, "Precondition Failed : Etag Mismatch", blobName)); throw new IOException("Unable to upload object [" + blobName + "] due to ETag mismatch", e); } else { IOException exception = new IOException( String.format(Locale.ROOT, "S3 error during multipart upload [%s]: %s", blobName, e.getMessage()), e ); - etagListener.onFailure(exception); + listener.onFailure(exception); throw exception; } } catch (SdkException e) { IOException exception = new IOException(String.format(Locale.ROOT, "S3 multipart upload failed for [%s]", blobName), e); - etagListener.onFailure(exception); + listener.onFailure(exception); throw exception; } catch (Exception e) { IOException exception = new IOException( String.format(Locale.ROOT, "Unexpected error during multipart upload [%s]: %s", blobName, e.getMessage()), e ); - etagListener.onFailure(exception); + listener.onFailure(exception); throw exception; } finally { if (!success && Strings.hasLength(uploadId.get())) { From 2f8f4de0c29d601302b3af2befc1b99b38429cb0 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Wed, 28 May 2025 13:16:37 +0530 Subject: [PATCH 10/20] Javadoc for builder --- .../org/opensearch/common/blobstore/ConditionalWrite.java | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/server/src/main/java/org/opensearch/common/blobstore/ConditionalWrite.java b/server/src/main/java/org/opensearch/common/blobstore/ConditionalWrite.java index 0b3f94aa603b8..2865c85e45ac5 100644 --- a/server/src/main/java/org/opensearch/common/blobstore/ConditionalWrite.java +++ b/server/src/main/java/org/opensearch/common/blobstore/ConditionalWrite.java @@ -84,6 +84,10 @@ public Instant getUnmodifiedSince() { return unmodifiedSince; } + /** + * Builder for {@link ConditionalWriteOptions}. + * Allows fine-grained construction of conditional write criteria. + */ public static final class Builder { private boolean ifNotExists = false; private boolean ifMatch = false; From 9fdefaa757dce30c41d2f252a6345597b3ff95b4 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Wed, 28 May 2025 13:18:01 +0530 Subject: [PATCH 11/20] asyncBlobUploadConditionally : Interfaces for AsyncMultiStreamBlobContainer & AsyncMultiStreamEncryptedBlobContainer --- .../AsyncMultiStreamBlobContainer.java | 20 +++++++++++++++++++ ...syncMultiStreamEncryptedBlobContainer.java | 10 ++++++++++ 2 files changed, 30 insertions(+) diff --git a/server/src/main/java/org/opensearch/common/blobstore/AsyncMultiStreamBlobContainer.java b/server/src/main/java/org/opensearch/common/blobstore/AsyncMultiStreamBlobContainer.java index b769cdc2fe7ab..08bdd90f617af 100644 --- a/server/src/main/java/org/opensearch/common/blobstore/AsyncMultiStreamBlobContainer.java +++ b/server/src/main/java/org/opensearch/common/blobstore/AsyncMultiStreamBlobContainer.java @@ -9,6 +9,8 @@ package org.opensearch.common.blobstore; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteOptions; +import org.opensearch.common.blobstore.ConditionalWrite.ConditionalWriteResponse; import org.opensearch.common.blobstore.stream.read.ReadContext; import org.opensearch.common.blobstore.stream.write.WriteContext; import org.opensearch.core.action.ActionListener; @@ -35,6 +37,24 @@ public interface AsyncMultiStreamBlobContainer extends BlobContainer { */ void asyncBlobUpload(WriteContext writeContext, ActionListener completionListener) throws IOException; + /** + * Reads blob content basis a preconditional requirement, from multiple streams each from a specific part of the file, which is provided by the + * StreamContextSupplier in the WriteContext passed to this method. An {@link IOException} is thrown if reading + * any of the input streams fails, or writing to the target blob fails + * + * @param writeContext A WriteContext object encapsulating all information needed to perform the upload + * @param options The {@link ConditionalWriteOptions} specifying the preconditions that must be met for the upload to proceed. + * @param completionListener The {@link ActionListener} to which upload events and the result will be published. + * @throws IOException if any of the input streams could not be read, or the target blob could not be written to + */ + default void asyncBlobUploadConditionally( + WriteContext writeContext, + ConditionalWriteOptions options, + ActionListener completionListener + ) throws IOException { + throw new UnsupportedOperationException("asyncBlobUploadConditionally is not implemented yet"); + }; + /** * Creates an async callback of a {@link ReadContext} containing the multipart streams for a specified blob within the container. * @param blobName The name of the blob for which the {@link ReadContext} needs to be fetched. diff --git a/server/src/main/java/org/opensearch/common/blobstore/AsyncMultiStreamEncryptedBlobContainer.java b/server/src/main/java/org/opensearch/common/blobstore/AsyncMultiStreamEncryptedBlobContainer.java index 286c01f9dca44..42019e49175ff 100644 --- a/server/src/main/java/org/opensearch/common/blobstore/AsyncMultiStreamEncryptedBlobContainer.java +++ b/server/src/main/java/org/opensearch/common/blobstore/AsyncMultiStreamEncryptedBlobContainer.java @@ -44,6 +44,16 @@ public void asyncBlobUpload(WriteContext writeContext, ActionListener comp blobContainer.asyncBlobUpload(encryptedWriteContext, completionListener); } + @Override + public void asyncBlobUploadConditionally( + WriteContext writeContext, + ConditionalWrite.ConditionalWriteOptions options, + ActionListener completionListener + ) throws IOException { + EncryptedWriteContext encryptedWriteContext = new EncryptedWriteContext<>(writeContext, cryptoHandler); + blobContainer.asyncBlobUploadConditionally(encryptedWriteContext, options, completionListener); + } + @Override public void readBlobAsync(String blobName, ActionListener listener) { try { From d8d8f9aa93d903ba103ac55c95599c768c36943c Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Wed, 28 May 2025 19:31:23 +0530 Subject: [PATCH 12/20] Log fix --- .../opensearch/repositories/s3/S3BlobContainer.java | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java index 1592ec1e31d53..1946496c898b0 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java @@ -256,9 +256,11 @@ public void asyncBlobUploadConditionally( ); } catch (Exception ex) { logger.error( - "Failed to conditionally upload large file {} of size {}", - uploadRequest.getKey(), - uploadRequest.getContentLength(), + () -> new ParameterizedMessage( + "Failed to conditionally upload large file {} of size {} ", + uploadRequest.getKey(), + uploadRequest.getContentLength() + ), ex ); completionListener.onFailure(ex); @@ -316,7 +318,7 @@ public void asyncBlobUploadConditionally( } } } catch (Exception e) { - logger.info("Exception during conditional blob upload for file {}", writeContext.getFileName(), e); + logger.info("Exception during conditional blob upload for file {}", writeContext.getFileName()); throw new IOException(e); } } From 5cfeb6e2fbf26e6f69873a2f10943899622b6227 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Thu, 29 May 2025 14:37:54 +0530 Subject: [PATCH 13/20] SpotlessApply post Conflict resolution --- .../repositories/s3/S3BlobContainer.java | 34 +++---- .../repositories/s3/async/UploadRequest.java | 2 +- .../s3/async/AsyncTransferManagerTests.java | 90 +++++++++++-------- 3 files changed, 70 insertions(+), 56 deletions(-) diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java index 417d60275a2e2..a8d2532b42fb2 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java @@ -233,10 +233,10 @@ public void asyncBlobUploadConditionally( writeContext.getExpectedChecksum(), blobStore.isUploadRetryEnabled(), writeContext.getMetadata(), - options, + options, blobStore.serverSideEncryptionType(), blobStore.serverSideEncryptionKmsKey(), - blobStore.serverSideEncryptionBucketKey(), + blobStore.serverSideEncryptionBucketKey(), blobStore.serverSideEncryptionEncryptionContext(), blobStore.expectedBucketOwner() ); @@ -338,21 +338,21 @@ public void asyncBlobUploadConditionally( @Override public void asyncBlobUpload(WriteContext writeContext, ActionListener completionListener) throws IOException { UploadRequest uploadRequest = new UploadRequest( - blobStore.bucket(), - buildKey(writeContext.getFileName()), - writeContext.getFileSize(), - writeContext.getWritePriority(), - writeContext.getUploadFinalizer(), - writeContext.doRemoteDataIntegrityCheck(), - writeContext.getExpectedChecksum(), - blobStore.isUploadRetryEnabled(), - writeContext.getMetadata(), - null, - blobStore.serverSideEncryptionType(), - blobStore.serverSideEncryptionKmsKey(), - blobStore.serverSideEncryptionBucketKey(), - blobStore.serverSideEncryptionEncryptionContext(), - blobStore.expectedBucketOwner() + blobStore.bucket(), + buildKey(writeContext.getFileName()), + writeContext.getFileSize(), + writeContext.getWritePriority(), + writeContext.getUploadFinalizer(), + writeContext.doRemoteDataIntegrityCheck(), + writeContext.getExpectedChecksum(), + blobStore.isUploadRetryEnabled(), + writeContext.getMetadata(), + null, + blobStore.serverSideEncryptionType(), + blobStore.serverSideEncryptionKmsKey(), + blobStore.serverSideEncryptionBucketKey(), + blobStore.serverSideEncryptionEncryptionContext(), + blobStore.expectedBucketOwner() ); try { // If file size is greater than the queue capacity than SizeBasedBlockingQ will always reject the upload. diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/UploadRequest.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/UploadRequest.java index c8d3cc599041d..f3b1e9271d781 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/UploadRequest.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/UploadRequest.java @@ -154,4 +154,4 @@ public String getServerSideEncryptionEncryptionContext() { public String getExpectedBucketOwner() { return expectedBucketOwner; } -} \ No newline at end of file +} diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java index 5ea2a31939f97..fd3489bcef6a2 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java @@ -99,10 +99,9 @@ public void testOneChunkUpload() { CompletableFuture resultFuture = asyncTransferManager.uploadObject( s3AsyncClient, new UploadRequest("bucket", "key", ByteSizeUnit.MB.toBytes(1), WritePriority.HIGH, uploadSuccess -> { - - }, false, null, true, metadata, null , ServerSideEncryption.AWS_KMS.toString(), randomAlphaOfLength(10), true, null, null), - + }, false, null, true, metadata, null, ServerSideEncryption.AWS_KMS.toString(), randomAlphaOfLength(10), true, null, null), + new StreamContext((partIdx, partSize, position) -> { streamRef.set(new ZeroInputStream(partSize)); return new InputStreamContainer(streamRef.get(), partSize, position); @@ -146,11 +145,26 @@ public void testOneChunkUploadCorruption() { Map metadata = new HashMap<>(); metadata.put("key1", "value1"); metadata.put("key2", "value2"); - + CompletableFuture resultFuture = asyncTransferManager.uploadObject( s3AsyncClient, - new UploadRequest("bucket", "key", ByteSizeUnit.MB.toBytes(1), WritePriority.HIGH, uploadSuccess -> { - }, false, null, true, metadata, null, ServerSideEncryption.AWS_KMS.toString(), randomAlphaOfLength(10), true, null, null), + new UploadRequest( + "bucket", + "key", + ByteSizeUnit.MB.toBytes(1), + WritePriority.HIGH, + uploadSuccess -> {}, + false, + null, + true, + metadata, + null, + ServerSideEncryption.AWS_KMS.toString(), + randomAlphaOfLength(10), + true, + null, + null + ), new StreamContext( (partIdx, partSize, position) -> new InputStreamContainer(new ZeroInputStream(partSize), partSize, position), ByteSizeUnit.MB.toBytes(1), @@ -206,7 +220,7 @@ public void testMultipartUpload() { CompletableFuture resultFuture = asyncTransferManager.uploadObject( s3AsyncClient, new UploadRequest("bucket", "key", ByteSizeUnit.MB.toBytes(5), WritePriority.HIGH, uploadSuccess -> { - + }, true, 3376132981L, true, metadata, null, ServerSideEncryption.AWS_KMS.toString(), randomAlphaOfLength(10), true, null, null), new StreamContext((partIdx, partSize, position) -> { @@ -323,12 +337,12 @@ public void testConditionalOneChunkUpload() { null, true, metadata, - options, - null, - null, - false, - null, - null + options, + null, + null, + false, + null, + null ), new StreamContext((partIdx, partSize, position) -> { streamRef.set(new ZeroInputStream(partSize)); @@ -381,12 +395,12 @@ public void testConditionalOneChunkUploadPreconditionFailed() { null, true, null, - options, - null, - null, - false, - null, - null + options, + null, + null, + false, + null, + null ), new StreamContext((partIdx, partSize, position) -> { streamRef.set(new ZeroInputStream(partSize)); @@ -451,12 +465,12 @@ public void testConditionalOneChunkUploadCorruption() { null, true, null, - options, - null, - null, - false, - null, - null + options, + null, + null, + false, + null, + null ), new StreamContext((partIdx, partSize, position) -> { streamRef.set(new ZeroInputStream(partSize)); @@ -528,12 +542,12 @@ public void testConditionalMultipartUploadPreconditionFailed() { 3376132981L, true, null, - options, - null, - null, - false, - null, - null + options, + null, + null, + false, + null, + null ), new StreamContext((partIdx, partSize, position) -> { InputStream stream = new ZeroInputStream(partSize); @@ -604,13 +618,13 @@ public void testConditionalMultipartUploadCorruption() { 0L, true, null, - options, - null, - null, - false, - null, - null - ), + options, + null, + null, + false, + null, + null + ), new StreamContext((partIdx, partSize, position) -> { InputStream stream = new ZeroInputStream(partSize); streams.add(stream); From 0a906807cd4948c7e3a3b53d2f20cf88092deed4 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Thu, 29 May 2025 15:06:00 +0530 Subject: [PATCH 14/20] Conflict resolution in executeMultipartUploadConditionally --- .../repositories/s3/S3BlobContainer.java | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java index a8d2532b42fb2..263fffea1c950 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java @@ -55,7 +55,6 @@ import software.amazon.awssdk.services.s3.model.ObjectAttributes; import software.amazon.awssdk.services.s3.model.PutObjectRequest; import software.amazon.awssdk.services.s3.model.S3Exception; -import software.amazon.awssdk.services.s3.model.ServerSideEncryption; import software.amazon.awssdk.services.s3.model.UploadPartRequest; import software.amazon.awssdk.services.s3.model.UploadPartResponse; import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Iterable; @@ -693,14 +692,13 @@ public void executeMultipartUploadConditionally( .key(blobName) .storageClass(blobStore.getStorageClass()) .acl(blobStore.getCannedACL()) - .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector)); + .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector)) + .expectedBucketOwner(blobStore.expectedBucketOwner()); - if (metadata != null && !metadata.isEmpty()) { + if (CollectionUtils.isNotEmpty(metadata)) { createRequestBuilder.metadata(metadata); } - if (blobStore.serverSideEncryption()) { - createRequestBuilder.serverSideEncryption(ServerSideEncryption.AES256); - } + configureEncryptionSettings(createRequestBuilder, blobStore); final CreateMultipartUploadRequest createMultipartUploadRequest = createRequestBuilder.build(); final SetOnce uploadId = new SetOnce<>(); @@ -733,6 +731,7 @@ public void executeMultipartUploadConditionally( .partNumber(i) .contentLength(currentPartSize) .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector)) + .expectedBucketOwner(blobStore.expectedBucketOwner()) .build(); bytesCount += currentPartSize; @@ -767,7 +766,8 @@ public void executeMultipartUploadConditionally( .key(blobName) .uploadId(uploadId.get()) .multipartUpload(CompletedMultipartUpload.builder().parts(parts).build()) - .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector)); + .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector)) + .expectedBucketOwner(blobStore.expectedBucketOwner()); if (options.isIfMatch()) { completeRequestBuilder.ifMatch(options.getVersionIdentifier()); @@ -821,6 +821,7 @@ public void executeMultipartUploadConditionally( .bucket(bucketName) .key(blobName) .uploadId(uploadId.get()) + .expectedBucketOwner(blobStore.expectedBucketOwner()) .build(); try (AmazonS3Reference abortClient = blobStore.clientReference()) { SocketAccess.doPrivilegedVoid(() -> abortClient.get().abortMultipartUpload(abortRequest)); From cbc018847651c6154509515421e06294426be931 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Thu, 29 May 2025 15:36:13 +0530 Subject: [PATCH 15/20] Conflict resolution in executeMultipartUploadConditionally's UTs --- .../repositories/s3/S3BlobStoreContainerTests.java | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java index d5a9533f37adf..1b0df7f2582d4 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java @@ -1022,9 +1022,6 @@ public void testExecuteMultipartUploadConditionallyWithEtagMatchSuccess() throws final StorageClass storageClass = randomFrom(StorageClass.values()); when(blobStore.getStorageClass()).thenReturn(storageClass); - final boolean serverSideEncryption = randomBoolean(); - when(blobStore.serverSideEncryption()).thenReturn(serverSideEncryption); - final ObjectCannedACL cannedAccessControlList = randomBoolean() ? randomFrom(ObjectCannedACL.values()) : null; if (cannedAccessControlList != null) { when(blobStore.getCannedACL()).thenReturn(cannedAccessControlList); @@ -1090,10 +1087,6 @@ public void testExecuteMultipartUploadConditionallyWithEtagMatchSuccess() throws assertEquals(cannedAccessControlList, createRequest.acl()); assertEquals(metadata, createRequest.metadata()); - if (serverSideEncryption) { - assertEquals(ServerSideEncryption.AES256, createRequest.serverSideEncryption()); - } - List partRequests = uploadPartRequestCaptor.getAllValues(); assertEquals(partCount, partRequests.size()); @@ -1154,9 +1147,6 @@ public void testExecuteMultipartUploadConditionallyWithMetadataAndSSE() throws I when(blobStore.bufferSizeInBytes()).thenReturn(partSize); when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); - final boolean serverSideEncryption = true; - when(blobStore.serverSideEncryption()).thenReturn(serverSideEncryption); - final StorageClass storageClass = randomFrom(StorageClass.values()); when(blobStore.getStorageClass()).thenReturn(storageClass); @@ -1283,7 +1273,6 @@ public void testExecuteMultipartUploadConditionallyContentIntegrity() throws IOE when(blobStore.bufferSizeInBytes()).thenReturn(partSize); when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); - when(blobStore.serverSideEncryption()).thenReturn(false); when(blobStore.isUploadRetryEnabled()).thenReturn(false); final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); @@ -1576,7 +1565,6 @@ public void testExecuteMultipartUploadConditionallyPreconditionFailed() { when(blobStore.bufferSizeInBytes()).thenReturn(partSize); when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); - when(blobStore.serverSideEncryption()).thenReturn(false); when(blobStore.isUploadRetryEnabled()).thenReturn(false); final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); @@ -1676,7 +1664,6 @@ public void testExecuteMultipartUploadConditionallyS3ExceptionTypes() { when(blobStore.bufferSizeInBytes()).thenReturn(partSize); when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); - when(blobStore.serverSideEncryption()).thenReturn(false); final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); From 0eedd2cdc061237b1bac19b089ac8ec605834ec1 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Thu, 29 May 2025 17:29:07 +0530 Subject: [PATCH 16/20] Conflict resolution in S3BlobContainerMockClientTests & AsyncTransferManagerTests --- .../s3/S3BlobContainerMockClientTests.java | 9 ++ .../s3/S3BlobStoreContainerTests.java | 131 +++++++++++++++++- .../s3/async/AsyncTransferManagerTests.java | 1 - 3 files changed, 138 insertions(+), 3 deletions(-) diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java index 4ba91a6467383..b77d6634a8995 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java @@ -841,6 +841,15 @@ private void testLargeFilesRedirectedToSlowSyncClientConditional( when(blobStore.getCannedACL()).thenReturn(cannedAccessControlList); } + if (randomBoolean()) { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AES256.toString()); + } else { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AWS_KMS.toString()); + when(blobStore.serverSideEncryptionKmsKey()).thenReturn(randomAlphaOfLength(10)); + when(blobStore.serverSideEncryptionBucketKey()).thenReturn(randomBoolean()); + when(blobStore.serverSideEncryptionEncryptionContext()).thenReturn(randomAlphaOfLength(10)); + } + final S3Client client = mock(S3Client.class); final AmazonS3Reference clientReference = Mockito.spy(new AmazonS3Reference(client)); doNothing().when(clientReference).close(); diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java index 1b0df7f2582d4..56c24fa43007d 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java @@ -1022,6 +1022,19 @@ public void testExecuteMultipartUploadConditionallyWithEtagMatchSuccess() throws final StorageClass storageClass = randomFrom(StorageClass.values()); when(blobStore.getStorageClass()).thenReturn(storageClass); + final boolean useSseKms = randomBoolean(); + final String kmsKeyId = randomAlphaOfLengthBetween(10, 20); + final String kmsContext = randomAlphaOfLengthBetween(10, 20); + final boolean useBucketKey = randomBoolean(); + if (useSseKms) { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AWS_KMS.toString()); + when(blobStore.serverSideEncryptionKmsKey()).thenReturn(kmsKeyId); + when(blobStore.serverSideEncryptionBucketKey()).thenReturn(useBucketKey); + when(blobStore.serverSideEncryptionEncryptionContext()).thenReturn(kmsContext); + } else { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AES256.toString()); + } + final ObjectCannedACL cannedAccessControlList = randomBoolean() ? randomFrom(ObjectCannedACL.values()) : null; if (cannedAccessControlList != null) { when(blobStore.getCannedACL()).thenReturn(cannedAccessControlList); @@ -1087,6 +1100,16 @@ public void testExecuteMultipartUploadConditionallyWithEtagMatchSuccess() throws assertEquals(cannedAccessControlList, createRequest.acl()); assertEquals(metadata, createRequest.metadata()); + // ENCRYPTION VERIFICATION: Updated encryption verification + if (useSseKms) { + assertEquals(ServerSideEncryption.AWS_KMS, createRequest.serverSideEncryption()); + assertEquals(kmsKeyId, createRequest.ssekmsKeyId()); + assertEquals(kmsContext, createRequest.ssekmsEncryptionContext()); + assertEquals(useBucketKey, createRequest.bucketKeyEnabled()); + } else { + assertEquals(ServerSideEncryption.AES256, createRequest.serverSideEncryption()); + } + List partRequests = uploadPartRequestCaptor.getAllValues(); assertEquals(partCount, partRequests.size()); @@ -1147,6 +1170,20 @@ public void testExecuteMultipartUploadConditionallyWithMetadataAndSSE() throws I when(blobStore.bufferSizeInBytes()).thenReturn(partSize); when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); + // ENCRYPTION CHANGES: Replace hard-coded encryption with enhanced configuration + final boolean useSseKms = randomBoolean(); + final String kmsKeyId = randomAlphaOfLengthBetween(10, 20); + final String kmsContext = randomAlphaOfLengthBetween(10, 20); + final boolean useBucketKey = randomBoolean(); + if (useSseKms) { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AWS_KMS.toString()); + when(blobStore.serverSideEncryptionKmsKey()).thenReturn(kmsKeyId); + when(blobStore.serverSideEncryptionBucketKey()).thenReturn(useBucketKey); + when(blobStore.serverSideEncryptionEncryptionContext()).thenReturn(kmsContext); + } else { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AES256.toString()); + } + final StorageClass storageClass = randomFrom(StorageClass.values()); when(blobStore.getStorageClass()).thenReturn(storageClass); @@ -1210,7 +1247,15 @@ public void testExecuteMultipartUploadConditionallyWithMetadataAndSSE() throws I assertEquals(cannedAccessControlList, createRequest.acl()); assertEquals(metadata, createRequest.metadata()); - assertEquals(ServerSideEncryption.AES256, createRequest.serverSideEncryption()); + // ENCRYPTION VERIFICATION: Updated encryption verification + if (useSseKms) { + assertEquals(ServerSideEncryption.AWS_KMS, createRequest.serverSideEncryption()); + assertEquals(kmsKeyId, createRequest.ssekmsKeyId()); + assertEquals(kmsContext, createRequest.ssekmsEncryptionContext()); + assertEquals(useBucketKey, createRequest.bucketKeyEnabled()); + } else { + assertEquals(ServerSideEncryption.AES256, createRequest.serverSideEncryption()); + } List partRequests = uploadPartRequestCaptor.getAllValues(); assertEquals(2, partRequests.size()); @@ -1273,6 +1318,21 @@ public void testExecuteMultipartUploadConditionallyContentIntegrity() throws IOE when(blobStore.bufferSizeInBytes()).thenReturn(partSize); when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); + + // ENCRYPTION CHANGES: Replace serverSideEncryption with enhanced configuration + final boolean useSseKms = randomBoolean(); + final String kmsKeyId = randomAlphaOfLengthBetween(10, 20); + final String kmsContext = randomAlphaOfLengthBetween(10, 20); + final boolean useBucketKey = randomBoolean(); + if (useSseKms) { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AWS_KMS.toString()); + when(blobStore.serverSideEncryptionKmsKey()).thenReturn(kmsKeyId); + when(blobStore.serverSideEncryptionBucketKey()).thenReturn(useBucketKey); + when(blobStore.serverSideEncryptionEncryptionContext()).thenReturn(kmsContext); + } else { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AES256.toString()); + } + when(blobStore.isUploadRetryEnabled()).thenReturn(false); final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); @@ -1328,6 +1388,8 @@ public void testExecuteMultipartUploadConditionallyContentIntegrity() throws IOE assertEquals(bucketName, createRequest.bucket()); assertEquals(blobPath.buildAsString() + blobName, createRequest.key()); + // No explicit encryption verification needed for this test as it focuses on content integrity + final CompleteMultipartUploadRequest completeRequest = completeRequestCaptor.getValue(); assertEquals(inputETag, completeRequest.ifMatch()); assertEquals(bucketName, completeRequest.bucket()); @@ -1417,6 +1479,19 @@ public void testExecuteMultipartUploadConditionallySizeValidation() { when(blobStore.bufferSizeInBytes()).thenReturn(ByteSizeUnit.MB.toBytes(5)); + final boolean useSseKms = randomBoolean(); + final String kmsKeyId = randomAlphaOfLengthBetween(10, 20); + final String kmsContext = randomAlphaOfLengthBetween(10, 20); + final boolean useBucketKey = randomBoolean(); + if (useSseKms) { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AWS_KMS.toString()); + when(blobStore.serverSideEncryptionKmsKey()).thenReturn(kmsKeyId); + when(blobStore.serverSideEncryptionBucketKey()).thenReturn(useBucketKey); + when(blobStore.serverSideEncryptionEncryptionContext()).thenReturn(kmsContext); + } else { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AES256.toString()); + } + ArgumentCaptor completeRequestCaptor = ArgumentCaptor.forClass( CompleteMultipartUploadRequest.class ); @@ -1565,6 +1640,20 @@ public void testExecuteMultipartUploadConditionallyPreconditionFailed() { when(blobStore.bufferSizeInBytes()).thenReturn(partSize); when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); + + final boolean useSseKms = randomBoolean(); + final String kmsKeyId = randomAlphaOfLengthBetween(10, 20); + final String kmsContext = randomAlphaOfLengthBetween(10, 20); + final boolean useBucketKey = randomBoolean(); + if (useSseKms) { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AWS_KMS.toString()); + when(blobStore.serverSideEncryptionKmsKey()).thenReturn(kmsKeyId); + when(blobStore.serverSideEncryptionBucketKey()).thenReturn(useBucketKey); + when(blobStore.serverSideEncryptionEncryptionContext()).thenReturn(kmsContext); + } else { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AES256.toString()); + } + when(blobStore.isUploadRetryEnabled()).thenReturn(false); final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); @@ -1665,6 +1754,19 @@ public void testExecuteMultipartUploadConditionallyS3ExceptionTypes() { when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); + final boolean useSseKms = randomBoolean(); + final String kmsKeyId = randomAlphaOfLengthBetween(10, 20); + final String kmsContext = randomAlphaOfLengthBetween(10, 20); + final boolean useBucketKey = randomBoolean(); + if (useSseKms) { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AWS_KMS.toString()); + when(blobStore.serverSideEncryptionKmsKey()).thenReturn(kmsKeyId); + when(blobStore.serverSideEncryptionBucketKey()).thenReturn(useBucketKey); + when(blobStore.serverSideEncryptionEncryptionContext()).thenReturn(kmsContext); + } else { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AES256.toString()); + } + final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); final S3Client client = mock(S3Client.class); @@ -1712,7 +1814,6 @@ public void testExecuteMultipartUploadConditionallyS3ExceptionTypes() { final AtomicReference capturedException = new AtomicReference<>(); ActionListener responseListener = ActionListener.wrap( - // Changed listener type r -> fail("Should have failed with exception"), capturedException::set ); @@ -1780,6 +1881,19 @@ public void testExecuteMultipartUploadConditionallySdkException() { when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); + final boolean useSseKms = randomBoolean(); + final String kmsKeyId = randomAlphaOfLengthBetween(10, 20); + final String kmsContext = randomAlphaOfLengthBetween(10, 20); + final boolean useBucketKey = randomBoolean(); + if (useSseKms) { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AWS_KMS.toString()); + when(blobStore.serverSideEncryptionKmsKey()).thenReturn(kmsKeyId); + when(blobStore.serverSideEncryptionBucketKey()).thenReturn(useBucketKey); + when(blobStore.serverSideEncryptionEncryptionContext()).thenReturn(kmsContext); + } else { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AES256.toString()); + } + final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); final S3Client client = mock(S3Client.class); @@ -1881,6 +1995,19 @@ enum ResourceScenario { when(blobStore.getStatsMetricPublisher()).thenReturn(metricPublisher); when(blobStore.getStorageClass()).thenReturn(StorageClass.STANDARD); + final boolean useSseKms = randomBoolean(); + final String kmsKeyId = randomAlphaOfLengthBetween(10, 20); + final String kmsContext = randomAlphaOfLengthBetween(10, 20); + final boolean useBucketKey = randomBoolean(); + if (useSseKms) { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AWS_KMS.toString()); + when(blobStore.serverSideEncryptionKmsKey()).thenReturn(kmsKeyId); + when(blobStore.serverSideEncryptionBucketKey()).thenReturn(useBucketKey); + when(blobStore.serverSideEncryptionEncryptionContext()).thenReturn(kmsContext); + } else { + when(blobStore.serverSideEncryptionType()).thenReturn(ServerSideEncryption.AES256.toString()); + } + final S3Client client = mock(S3Client.class); AmazonS3Reference primaryClientReference = mock(AmazonS3Reference.class); diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java index fd3489bcef6a2..9928211c54479 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java @@ -641,7 +641,6 @@ public void testConditionalMultipartUploadCorruption() { assertNotNull("Exception should be a CorruptFileException", throwable); assertTrue("Exception should be a CorruptFileException", throwable instanceof CorruptFileException); } - verify(s3AsyncClient, times(1)).createMultipartUpload(any(CreateMultipartUploadRequest.class)); verify(s3AsyncClient, times(5)).uploadPart(any(UploadPartRequest.class), any(AsyncRequestBody.class)); verify(s3AsyncClient, times(0)).completeMultipartUpload(any(CompleteMultipartUploadRequest.class)); From 0e31c6270182f4abdbdf32c15dd3ee1ecbe31596 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Wed, 11 Jun 2025 17:01:16 +0530 Subject: [PATCH 17/20] S3BlobContainer Refactor --- .../repositories/s3/S3BlobContainer.java | 218 +++++++----------- 1 file changed, 82 insertions(+), 136 deletions(-) diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java index 263fffea1c950..f058f458989d5 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java @@ -221,6 +221,19 @@ public void asyncBlobUploadConditionally( WriteContext writeContext, ConditionalWriteOptions options, ActionListener completionListener + ) throws IOException { + executeAsyncUpload(writeContext, options, completionListener); + } + + @Override + public void asyncBlobUpload(WriteContext writeContext, ActionListener completionListener) throws IOException { + executeAsyncUpload(writeContext, null, new ConditionalResponseToVoidListener(completionListener)); + } + + private void executeAsyncUpload( + WriteContext writeContext, + @Nullable ConditionalWriteOptions options, + ActionListener completionListener ) throws IOException { UploadRequest uploadRequest = new UploadRequest( blobStore.bucket(), @@ -256,125 +269,26 @@ public void asyncBlobUploadConditionally( InputStreamContainer inputStream = streamContext.provideStream(0); try { - executeMultipartUploadConditionally( - blobStore, - buildKey(writeContext.getFileName()), - inputStream.getInputStream(), - uploadRequest.getContentLength(), - uploadRequest.getMetadata(), - options, - completionListener - ); - } catch (Exception ex) { - logger.error( - () -> new ParameterizedMessage( - "Failed to conditionally upload large file {} of size {} ", + if (options != null) { + executeMultipartUploadConditionally( + blobStore, uploadRequest.getKey(), - uploadRequest.getContentLength() - ), - ex - ); - completionListener.onFailure(ex); - } - return; - } - - long partSize = blobStore.getAsyncTransferManager() - .calculateOptimalPartSize(writeContext.getFileSize(), writeContext.getWritePriority(), blobStore.isUploadRetryEnabled()); - - StreamContext streamContext = SocketAccess.doPrivileged(() -> writeContext.getStreamProvider(partSize)); - - try (AmazonAsyncS3Reference amazonS3Reference = SocketAccess.doPrivileged(blobStore::asyncClientReference)) { - S3AsyncClient s3AsyncClient; - if (writeContext.getWritePriority() == WritePriority.URGENT) { - s3AsyncClient = amazonS3Reference.get().urgentClient(); - } else if (writeContext.getWritePriority() == WritePriority.HIGH) { - s3AsyncClient = amazonS3Reference.get().priorityClient(); - } else { - s3AsyncClient = amazonS3Reference.get().client(); - } - - if (writeContext.getWritePriority() == WritePriority.URGENT - || writeContext.getWritePriority() == WritePriority.HIGH - || blobStore.isPermitBackedTransferEnabled() == false) { - createFileCompletableFutureConditionally(s3AsyncClient, uploadRequest, streamContext, completionListener); - } else if (writeContext.getWritePriority() == WritePriority.LOW) { - blobStore.getLowPrioritySizeBasedBlockingQ() - .produce( - new SizeBasedBlockingQ.Item( - writeContext.getFileSize(), - () -> createFileCompletableFutureConditionally( - s3AsyncClient, - uploadRequest, - streamContext, - completionListener - ) - ) + inputStream.getInputStream(), + uploadRequest.getContentLength(), + uploadRequest.getMetadata(), + options, + completionListener ); - } else if (writeContext.getWritePriority() == WritePriority.NORMAL) { - blobStore.getNormalPrioritySizeBasedBlockingQ() - .produce( - new SizeBasedBlockingQ.Item( - writeContext.getFileSize(), - () -> createFileCompletableFutureConditionally( - s3AsyncClient, - uploadRequest, - streamContext, - completionListener - ) - ) + } else { + executeMultipartUpload( + blobStore, + uploadRequest.getKey(), + inputStream.getInputStream(), + uploadRequest.getContentLength(), + uploadRequest.getMetadata() ); - } else { - throw new IllegalStateException("Cannot perform upload for other priority types."); - } - } - } catch (Exception e) { - logger.info("Exception during conditional blob upload for file {}", writeContext.getFileName()); - throw new IOException(e); - } - } - - @Override - public void asyncBlobUpload(WriteContext writeContext, ActionListener completionListener) throws IOException { - UploadRequest uploadRequest = new UploadRequest( - blobStore.bucket(), - buildKey(writeContext.getFileName()), - writeContext.getFileSize(), - writeContext.getWritePriority(), - writeContext.getUploadFinalizer(), - writeContext.doRemoteDataIntegrityCheck(), - writeContext.getExpectedChecksum(), - blobStore.isUploadRetryEnabled(), - writeContext.getMetadata(), - null, - blobStore.serverSideEncryptionType(), - blobStore.serverSideEncryptionKmsKey(), - blobStore.serverSideEncryptionBucketKey(), - blobStore.serverSideEncryptionEncryptionContext(), - blobStore.expectedBucketOwner() - ); - try { - // If file size is greater than the queue capacity than SizeBasedBlockingQ will always reject the upload. - // Therefore, redirecting it to slow client. - if ((uploadRequest.getWritePriority() == WritePriority.LOW - && blobStore.getLowPrioritySizeBasedBlockingQ().isMaxCapacityBelowContentLength(uploadRequest.getContentLength()) == false) - || (uploadRequest.getWritePriority() != WritePriority.HIGH - && uploadRequest.getWritePriority() != WritePriority.URGENT - && blobStore.getNormalPrioritySizeBasedBlockingQ() - .isMaxCapacityBelowContentLength(uploadRequest.getContentLength()) == false)) { - StreamContext streamContext = SocketAccess.doPrivileged( - () -> writeContext.getStreamProvider(uploadRequest.getContentLength()) - ); - InputStreamContainer inputStream = streamContext.provideStream(0); - try { - executeMultipartUpload( - blobStore, - uploadRequest.getKey(), - inputStream.getInputStream(), - uploadRequest.getContentLength(), - uploadRequest.getMetadata() - ); - completionListener.onResponse(null); + completionListener.onResponse(ConditionalWriteResponse.success(null)); + } } catch (Exception ex) { logger.error( () -> new ParameterizedMessage( @@ -405,23 +319,29 @@ public void asyncBlobUpload(WriteContext writeContext, ActionListener comp if (writeContext.getWritePriority() == WritePriority.URGENT || writeContext.getWritePriority() == WritePriority.HIGH || blobStore.isPermitBackedTransferEnabled() == false) { - createFileCompletableFuture(s3AsyncClient, uploadRequest, streamContext, completionListener); + + if (options != null) { + createFileCompletableFutureConditionally(s3AsyncClient, uploadRequest, streamContext, completionListener); + } else { + createFileCompletableFuture(s3AsyncClient, uploadRequest, streamContext, completionListener); + } + } else if (writeContext.getWritePriority() == WritePriority.LOW) { - blobStore.getLowPrioritySizeBasedBlockingQ() - .produce( - new SizeBasedBlockingQ.Item( - writeContext.getFileSize(), - () -> createFileCompletableFuture(s3AsyncClient, uploadRequest, streamContext, completionListener) - ) - ); + blobStore.getLowPrioritySizeBasedBlockingQ().produce(new SizeBasedBlockingQ.Item(writeContext.getFileSize(), () -> { + if (options != null) { + createFileCompletableFutureConditionally(s3AsyncClient, uploadRequest, streamContext, completionListener); + } else { + createFileCompletableFuture(s3AsyncClient, uploadRequest, streamContext, completionListener); + } + })); } else if (writeContext.getWritePriority() == WritePriority.NORMAL) { - blobStore.getNormalPrioritySizeBasedBlockingQ() - .produce( - new SizeBasedBlockingQ.Item( - writeContext.getFileSize(), - () -> createFileCompletableFuture(s3AsyncClient, uploadRequest, streamContext, completionListener) - ) - ); + blobStore.getNormalPrioritySizeBasedBlockingQ().produce(new SizeBasedBlockingQ.Item(writeContext.getFileSize(), () -> { + if (options != null) { + createFileCompletableFutureConditionally(s3AsyncClient, uploadRequest, streamContext, completionListener); + } else { + createFileCompletableFuture(s3AsyncClient, uploadRequest, streamContext, completionListener); + } + })); } else { throw new IllegalStateException("Cannot perform upload for other priority types."); } @@ -432,15 +352,20 @@ public void asyncBlobUpload(WriteContext writeContext, ActionListener comp } } - private CompletableFuture createFileCompletableFuture( + private CompletableFuture createFileCompletableFuture( S3AsyncClient s3AsyncClient, UploadRequest uploadRequest, StreamContext streamContext, - ActionListener completionListener + ActionListener completionListener ) { - CompletableFuture completableFuture = blobStore.getAsyncTransferManager() + CompletableFuture standardFuture = blobStore.getAsyncTransferManager() .uploadObject(s3AsyncClient, uploadRequest, streamContext, blobStore.getStatsMetricPublisher()); - return completableFuture.whenComplete((response, throwable) -> { + + CompletableFuture convertedFuture = standardFuture.thenApply( + result -> ConditionalWriteResponse.success(null) + ); + + return convertedFuture.whenComplete((response, throwable) -> { if (throwable == null) { completionListener.onResponse(response); } else { @@ -469,6 +394,27 @@ private CompletableFuture createFileCompletableFutureC }); } + /** + * Helper class to convert ConditionalWriteResponse to Void + */ + private static class ConditionalResponseToVoidListener implements ActionListener { + private final ActionListener delegate; + + ConditionalResponseToVoidListener(ActionListener delegate) { + this.delegate = delegate; + } + + @Override + public void onResponse(ConditionalWriteResponse response) { + delegate.onResponse(null); + } + + @Override + public void onFailure(Exception e) { + delegate.onFailure(e); + } + } + @ExperimentalApi @Override public void readBlobAsync(String blobName, ActionListener listener) { @@ -794,7 +740,7 @@ public void executeMultipartUploadConditionally( } catch (S3Exception e) { if (e.statusCode() == HTTP_STATUS_PRECONDITION_FAILED) { - listener.onFailure(new OpenSearchException("stale_primary_shard", e, "Precondition Failed : Etag Mismatch", blobName)); + listener.onFailure(new OpenSearchException("Precondition Failed : Etag Mismatch", e, blobName)); throw new IOException("Unable to upload object [" + blobName + "] due to ETag mismatch", e); } else { IOException exception = new IOException( From 50498ac63f93ffbc1af8f51c4ffd2aa9761a4f92 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Wed, 11 Jun 2025 17:01:27 +0530 Subject: [PATCH 18/20] AsyncTransferManager Refactor --- .../s3/async/AsyncTransferManager.java | 603 ++++++------------ 1 file changed, 210 insertions(+), 393 deletions(-) diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java index 482d6ff712b1d..94bbeeb90ceca 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java @@ -14,7 +14,6 @@ import software.amazon.awssdk.services.s3.S3AsyncClient; import software.amazon.awssdk.services.s3.model.ChecksumAlgorithm; import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest; -import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse; import software.amazon.awssdk.services.s3.model.CompletedMultipartUpload; import software.amazon.awssdk.services.s3.model.CompletedPart; import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest; @@ -53,7 +52,6 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Semaphore; import java.util.concurrent.atomic.AtomicReferenceArray; -import java.util.function.BiFunction; import java.util.function.Supplier; import java.util.stream.IntStream; @@ -119,35 +117,7 @@ public CompletableFuture uploadObject( StreamContext streamContext, StatsMetricPublisher statsMetricPublisher ) { - - CompletableFuture returnFuture = new CompletableFuture<>(); - try { - if (streamContext.getNumberOfParts() == 1) { - log.debug(() -> "Starting the upload as a single upload part request"); - TransferSemaphoresHolder.RequestContext requestContext = transferSemaphoresHolder.createRequestContext(); - Semaphore semaphore = AsyncPartsHandler.maybeAcquireSemaphore( - transferSemaphoresHolder, - requestContext, - uploadRequest.getWritePriority(), - uploadRequest.getKey() - ); - try { - uploadInOneChunk(s3AsyncClient, uploadRequest, streamContext, returnFuture, statsMetricPublisher, semaphore); - } catch (Exception ex) { - if (semaphore != null) { - semaphore.release(); - } - throw ex; - } - } else { - log.debug(() -> "Starting the upload as multipart upload request"); - uploadInParts(s3AsyncClient, uploadRequest, streamContext, returnFuture, statsMetricPublisher); - } - } catch (Throwable throwable) { - returnFuture.completeExceptionally(throwable); - } - - return returnFuture; + return processUploadRequest(s3AsyncClient, uploadRequest, streamContext, statsMetricPublisher, false).thenApply(response -> null); } /** @@ -170,10 +140,25 @@ public CompletableFuture uploadObjectConditionally( throw new IllegalArgumentException("Cannot perform conditional upload with null options"); } + return processUploadRequest(s3AsyncClient, uploadRequest, streamContext, statsMetricPublisher, true); + } + + /** + internal upload execution for both regular and conditional uploads + */ + private CompletableFuture processUploadRequest( + S3AsyncClient s3AsyncClient, + UploadRequest uploadRequest, + StreamContext streamContext, + StatsMetricPublisher statsMetricPublisher, + boolean isConditional + ) { CompletableFuture returnFuture = new CompletableFuture<>(); try { if (streamContext.getNumberOfParts() == 1) { - log.debug(() -> "Starting conditional single part upload for key: " + uploadRequest.getKey()); + log.debug( + () -> "Starting " + (isConditional ? "conditional " : "") + "single part upload for key: " + uploadRequest.getKey() + ); TransferSemaphoresHolder.RequestContext requestContext = transferSemaphoresHolder.createRequestContext(); Semaphore semaphore = AsyncPartsHandler.maybeAcquireSemaphore( transferSemaphoresHolder, @@ -182,13 +167,14 @@ public CompletableFuture uploadObjectConditionally( uploadRequest.getKey() ); try { - uploadInOneChunkConditionally( + uploadInOneChunk( s3AsyncClient, uploadRequest, streamContext, returnFuture, statsMetricPublisher, - semaphore + semaphore, + isConditional ); } catch (Exception ex) { if (semaphore != null) { @@ -197,8 +183,10 @@ public CompletableFuture uploadObjectConditionally( throw ex; } } else { - log.debug(() -> "Starting conditional multipart upload for key: " + uploadRequest.getKey()); - uploadInPartsConditionally(s3AsyncClient, uploadRequest, streamContext, returnFuture, statsMetricPublisher); + log.debug( + () -> "Starting " + (isConditional ? "conditional " : "") + "multipart upload for key: " + uploadRequest.getKey() + ); + uploadInParts(s3AsyncClient, uploadRequest, streamContext, returnFuture, statsMetricPublisher, isConditional); } } catch (Throwable throwable) { returnFuture.completeExceptionally(throwable); @@ -207,12 +195,16 @@ public CompletableFuture uploadObjectConditionally( return returnFuture; } + /** + multipart upload initiation for both regular and conditional async uploads + */ private void uploadInParts( S3AsyncClient s3AsyncClient, UploadRequest uploadRequest, StreamContext streamContext, - CompletableFuture returnFuture, - StatsMetricPublisher statsMetricPublisher + CompletableFuture returnFuture, + StatsMetricPublisher statsMetricPublisher, + boolean isConditional ) { CreateMultipartUploadRequest.Builder createMultipartUploadRequestBuilder = CreateMultipartUploadRequest.builder() @@ -245,121 +237,24 @@ private void uploadInParts( uploadId = createMultipartUploadResponse.uploadId(); log.debug(() -> "Initiated new multipart upload, uploadId: " + createMultipartUploadResponse.uploadId()); } catch (Exception ex) { - handleException(returnFuture, () -> "Failed to initiate multipart upload", ex); + handleUploadException(returnFuture, ex, uploadRequest.getKey(), isConditional); return; } - doUploadInParts(s3AsyncClient, uploadRequest, streamContext, returnFuture, uploadId, statsMetricPublisher); - } - - private void uploadInPartsConditionally( - S3AsyncClient s3AsyncClient, - UploadRequest uploadRequest, - StreamContext streamContext, - CompletableFuture returnFuture, - StatsMetricPublisher statsMetricPublisher - ) { - CreateMultipartUploadRequest.Builder createMultipartUploadRequestBuilder = CreateMultipartUploadRequest.builder() - .bucket(uploadRequest.getBucket()) - .key(uploadRequest.getKey()) - .overrideConfiguration(o -> o.addMetricPublisher(statsMetricPublisher.multipartUploadMetricCollector)); - - if (CollectionUtils.isNotEmpty(uploadRequest.getMetadata())) { - createMultipartUploadRequestBuilder.metadata(uploadRequest.getMetadata()); - } - if (uploadRequest.doRemoteDataIntegrityCheck()) { - createMultipartUploadRequestBuilder.checksumAlgorithm(ChecksumAlgorithm.CRC32); - } - - CompletableFuture createMultipartUploadFuture = SocketAccess.doPrivileged( - () -> s3AsyncClient.createMultipartUpload(createMultipartUploadRequestBuilder.build()) - ); - - // Ensure cancellations are forwarded to the createMultipartUploadFuture future - CompletableFutureUtils.forwardExceptionTo(returnFuture, createMultipartUploadFuture); - - String uploadId; - try { - // Block main thread here so that upload of parts doesn't get executed in future completion thread. - // We should never execute latent operation like acquisition of permit in future completion pool. - CreateMultipartUploadResponse createMultipartUploadResponse = createMultipartUploadFuture.get(); - uploadId = createMultipartUploadResponse.uploadId(); - log.debug(() -> "Initiated new multipart upload, uploadId: " + createMultipartUploadResponse.uploadId()); - } catch (Exception ex) { - handleConditionalException(returnFuture, ex, uploadRequest.getKey()); - return; - } - - doUploadInPartsConditionally(s3AsyncClient, uploadRequest, streamContext, returnFuture, uploadId, statsMetricPublisher); + doUploadInParts(s3AsyncClient, uploadRequest, streamContext, returnFuture, uploadId, statsMetricPublisher, isConditional); } + /** + multipart upload execution for both regular and conditional uploads + */ private void doUploadInParts( - S3AsyncClient s3AsyncClient, - UploadRequest uploadRequest, - StreamContext streamContext, - CompletableFuture returnFuture, - String uploadId, - StatsMetricPublisher statsMetricPublisher - ) { - - // The list of completed parts must be sorted - AtomicReferenceArray completedParts = new AtomicReferenceArray<>(streamContext.getNumberOfParts()); - AtomicReferenceArray inputStreamContainers = new AtomicReferenceArray<>(streamContext.getNumberOfParts()); - - List> futures; - try { - futures = AsyncPartsHandler.uploadParts( - s3AsyncClient, - executorService, - priorityExecutorService, - urgentExecutorService, - uploadRequest, - streamContext, - uploadId, - completedParts, - inputStreamContainers, - statsMetricPublisher, - uploadRequest.isUploadRetryEnabled(), - transferSemaphoresHolder, - maxRetryablePartSize - ); - } catch (Exception ex) { - try { - AsyncPartsHandler.cleanUpParts(s3AsyncClient, uploadRequest, uploadId); - } finally { - returnFuture.completeExceptionally(ex); - } - return; - } - - CompletableFutureUtils.allOfExceptionForwarded(futures.toArray(CompletableFuture[]::new)).thenApply(resp -> { - try { - uploadRequest.getUploadFinalizer().accept(true); - } catch (IOException e) { - throw new RuntimeException(e); - } - return resp; - }).thenApply(ignore -> { - if (uploadRequest.doRemoteDataIntegrityCheck()) { - mergeAndVerifyChecksum(inputStreamContainers, uploadRequest.getKey(), uploadRequest.getExpectedChecksum()); - } - return null; - }) - .thenCompose(ignore -> completeMultipartUpload(s3AsyncClient, uploadRequest, uploadId, completedParts, statsMetricPublisher)) - .handle(handleExceptionOrResponse(s3AsyncClient, uploadRequest, returnFuture, uploadId)) - .exceptionally(throwable -> { - handleException(returnFuture, () -> "Unexpected exception occurred", throwable); - return null; - }); - } - - private void doUploadInPartsConditionally( S3AsyncClient s3AsyncClient, UploadRequest uploadRequest, StreamContext streamContext, CompletableFuture returnFuture, String uploadId, - StatsMetricPublisher statsMetricPublisher + StatsMetricPublisher statsMetricPublisher, + boolean isConditional ) { // The list of completed parts must be sorted AtomicReferenceArray completedParts = new AtomicReferenceArray<>(streamContext.getNumberOfParts()); @@ -404,207 +299,96 @@ private void doUploadInPartsConditionally( } return null; }).thenCompose(ignore -> { - log.debug(() -> "Completing conditional multipart upload, uploadId: " + uploadId); - - CompletedPart[] parts = IntStream.range(0, completedParts.length()).mapToObj(completedParts::get).toArray(CompletedPart[]::new); - - CompleteMultipartUploadRequest.Builder completeRequestBuilder = CompleteMultipartUploadRequest.builder() - .bucket(uploadRequest.getBucket()) - .key(uploadRequest.getKey()) - .uploadId(uploadId) - .multipartUpload(CompletedMultipartUpload.builder().parts(parts).build()) - .overrideConfiguration(o -> o.addMetricPublisher(statsMetricPublisher.multipartUploadMetricCollector)); - - if (uploadRequest.getConditionalOptions() != null) { - applyConditionalHeaders(completeRequestBuilder, uploadRequest.getConditionalOptions()); - } - - return s3AsyncClient.completeMultipartUpload(completeRequestBuilder.build()); + log.debug(() -> "Completing " + (isConditional ? "conditional " : "") + "multipart upload, uploadId: " + uploadId); + return completeMultipartUpload(s3AsyncClient, uploadRequest, uploadId, completedParts, statsMetricPublisher, isConditional); }).handle((response, throwable) -> { if (throwable != null) { AsyncPartsHandler.cleanUpParts(s3AsyncClient, uploadRequest, uploadId); - Throwable unwrappedThrowable = ExceptionsHelper.unwrap(throwable, S3Exception.class); - if (unwrappedThrowable != null) { - S3Exception s3Exception = (S3Exception) unwrappedThrowable; - if (s3Exception.statusCode() == HTTP_STATUS_PRECONDITION_FAILED) { - returnFuture.completeExceptionally( - S3Exception.builder() - .message("Conditional write failed: condition not met for " + uploadRequest.getKey()) - .statusCode(HTTP_STATUS_PRECONDITION_FAILED) - .cause(s3Exception) - .build() - ); - return null; - } else if (s3Exception.statusCode() == HTTP_STATUS_CONFLICT) { - returnFuture.completeExceptionally( - S3Exception.builder() - .message("Blob already exists: " + uploadRequest.getKey()) - .statusCode(HTTP_STATUS_CONFLICT) - .cause(s3Exception) - .build() - ); - return null; + if (isConditional) { + Throwable unwrappedThrowable = ExceptionsHelper.unwrap(throwable, S3Exception.class); + if (unwrappedThrowable != null) { + S3Exception s3Exception = (S3Exception) unwrappedThrowable; + if (s3Exception.statusCode() == HTTP_STATUS_PRECONDITION_FAILED) { + returnFuture.completeExceptionally( + S3Exception.builder() + .message("Conditional write failed: condition not met for " + uploadRequest.getKey()) + .statusCode(HTTP_STATUS_PRECONDITION_FAILED) + .cause(s3Exception) + .build() + ); + return null; + } else if (s3Exception.statusCode() == HTTP_STATUS_CONFLICT) { + returnFuture.completeExceptionally( + S3Exception.builder() + .message("Blob already exists: " + uploadRequest.getKey()) + .statusCode(HTTP_STATUS_CONFLICT) + .cause(s3Exception) + .build() + ); + return null; + } } } - handleConditionalException(returnFuture, throwable, uploadRequest.getKey()); + handleUploadException(returnFuture, throwable, uploadRequest.getKey(), isConditional); return null; } else { - returnFuture.complete(ConditionalWriteResponse.success(response.eTag())); + returnFuture.complete(response); return null; } + }).exceptionally(throwable -> { - handleConditionalException(returnFuture, throwable, uploadRequest.getKey()); + handleUploadException(returnFuture, throwable, uploadRequest.getKey(), isConditional); return null; }); } - private void mergeAndVerifyChecksum( - AtomicReferenceArray inputStreamContainers, - String fileName, - long expectedChecksum - ) { - long resultantChecksum = fromBase64String(inputStreamContainers.get(0).getChecksum()); - for (int index = 1; index < inputStreamContainers.length(); index++) { - long curChecksum = fromBase64String(inputStreamContainers.get(index).getChecksum()); - resultantChecksum = JZlib.crc32_combine(resultantChecksum, curChecksum, inputStreamContainers.get(index).getContentLength()); - } - - if (resultantChecksum != expectedChecksum) { - throw new RuntimeException(new CorruptFileException("File level checksums didn't match combined part checksums", fileName)); - } - } - - private BiFunction handleExceptionOrResponse( - S3AsyncClient s3AsyncClient, - UploadRequest uploadRequest, - CompletableFuture returnFuture, - String uploadId - ) { - - return (response, throwable) -> { - if (throwable != null) { - AsyncPartsHandler.cleanUpParts(s3AsyncClient, uploadRequest, uploadId); - handleException(returnFuture, () -> "Failed to send multipart upload requests.", throwable); - } else { - returnFuture.complete(null); - } - - return null; - }; - } - - private CompletableFuture completeMultipartUpload( + /** + complete multipart upload for both regular and conditional uploads + */ + private CompletableFuture completeMultipartUpload( S3AsyncClient s3AsyncClient, UploadRequest uploadRequest, String uploadId, AtomicReferenceArray completedParts, - StatsMetricPublisher statsMetricPublisher + StatsMetricPublisher statsMetricPublisher, + boolean isConditional ) { log.debug(() -> new ParameterizedMessage("Sending completeMultipartUploadRequest, uploadId: {}", uploadId)); CompletedPart[] parts = IntStream.range(0, completedParts.length()).mapToObj(completedParts::get).toArray(CompletedPart[]::new); - CompleteMultipartUploadRequest completeMultipartUploadRequest = CompleteMultipartUploadRequest.builder() + + CompleteMultipartUploadRequest.Builder completeRequestBuilder = CompleteMultipartUploadRequest.builder() .bucket(uploadRequest.getBucket()) .key(uploadRequest.getKey()) .uploadId(uploadId) .overrideConfiguration(o -> o.addMetricPublisher(statsMetricPublisher.multipartUploadMetricCollector)) .multipartUpload(CompletedMultipartUpload.builder().parts(parts).build()) - .expectedBucketOwner(uploadRequest.getExpectedBucketOwner()) - .build(); - - return SocketAccess.doPrivileged(() -> s3AsyncClient.completeMultipartUpload(completeMultipartUploadRequest)); - } - - private static String base64StringFromLong(Long val) { - return Base64.getEncoder().encodeToString(Arrays.copyOfRange(ByteUtils.toByteArrayBE(val), 4, 8)); - } + .expectedBucketOwner(uploadRequest.getExpectedBucketOwner()); - private static long fromBase64String(String base64String) { - byte[] decodedBytes = Base64.getDecoder().decode(base64String); - if (decodedBytes.length != 4) { - throw new IllegalArgumentException("Invalid Base64 encoded CRC32 checksum"); + if (isConditional && uploadRequest.getConditionalOptions() != null) { + applyConditionalHeaders(completeRequestBuilder, uploadRequest.getConditionalOptions()); } - long result = 0; - for (int i = 0; i < 4; i++) { - result <<= 8; - result |= (decodedBytes[i] & 0xFF); - } - return result; - } - private static void handleException(CompletableFuture returnFuture, Supplier message, Throwable throwable) { - Throwable cause = throwable instanceof CompletionException ? throwable.getCause() : throwable; + CompleteMultipartUploadRequest completeRequest = completeRequestBuilder.build(); - if (cause instanceof Error) { - returnFuture.completeExceptionally(cause); - } else { - SdkClientException exception = SdkClientException.create(message.get(), cause); - returnFuture.completeExceptionally(exception); - } + return SocketAccess.doPrivileged(() -> s3AsyncClient.completeMultipartUpload(completeRequest)) + .thenApply(response -> ConditionalWriteResponse.success(response.eTag())); } /** - * Error handler for conditional uploads - */ - private void handleConditionalException(CompletableFuture returnFuture, Throwable throwable, String resourceName) { - Throwable cause = throwable; - while (cause instanceof CompletionException && cause.getCause() != null) { - cause = cause.getCause(); - } - if (cause instanceof S3Exception s3e) { - if (s3e.statusCode() == HTTP_STATUS_PRECONDITION_FAILED) { - returnFuture.completeExceptionally( - S3Exception.builder() - .message("Conditional write failed: condition not met for " + resourceName) - .statusCode(HTTP_STATUS_PRECONDITION_FAILED) - .cause(s3e) - .build() - ); - return; - } else if (s3e.statusCode() == HTTP_STATUS_CONFLICT) { - returnFuture.completeExceptionally( - S3Exception.builder() - .message("Blob already exists: " + resourceName) - .statusCode(HTTP_STATUS_CONFLICT) - .cause(s3e) - .build() - ); - return; - } - } - if (cause instanceof Error) { - returnFuture.completeExceptionally(cause); - } else { - SdkClientException exception = SdkClientException.create("Failed conditional upload of " + resourceName, cause); - returnFuture.completeExceptionally(exception); - } - } - - /** - * Calculates the optimal part size of each part request if the upload operation is carried out as multipart upload. + single chunk upload for both regular and conditional uploads */ - public long calculateOptimalPartSize(long contentLengthOfSource, WritePriority writePriority, boolean uploadRetryEnabled) { - if (contentLengthOfSource < ByteSizeUnit.MB.toBytes(5)) { - return contentLengthOfSource; - } - if (uploadRetryEnabled && (writePriority == WritePriority.HIGH || writePriority == WritePriority.URGENT)) { - return new ByteSizeValue(5, ByteSizeUnit.MB).getBytes(); - } - double optimalPartSize = contentLengthOfSource / (double) MAX_UPLOAD_PARTS; - optimalPartSize = Math.ceil(optimalPartSize); - return (long) Math.max(optimalPartSize, minimumPartSize); - } - @SuppressWarnings("unchecked") private void uploadInOneChunk( S3AsyncClient s3AsyncClient, UploadRequest uploadRequest, StreamContext streamContext, - CompletableFuture returnFuture, + CompletableFuture returnFuture, StatsMetricPublisher statsMetricPublisher, - Semaphore semaphore + Semaphore semaphore, + boolean isConditional ) { PutObjectRequest.Builder putObjectRequestBuilder = PutObjectRequest.builder() .bucket(uploadRequest.getBucket()) @@ -623,6 +407,10 @@ private void uploadInOneChunk( configureEncryptionSettings(putObjectRequestBuilder, uploadRequest); + if (isConditional && uploadRequest.getConditionalOptions() != null) { + applyConditionalHeaders(putObjectRequestBuilder, uploadRequest.getConditionalOptions()); + } + PutObjectRequest putObjectRequest = putObjectRequestBuilder.build(); ExecutorService streamReadExecutor; if (uploadRequest.getWritePriority() == WritePriority.URGENT) { @@ -633,7 +421,7 @@ private void uploadInOneChunk( streamReadExecutor = executorService; } - CompletableFuture putObjectFuture = SocketAccess.doPrivileged(() -> { + CompletableFuture putObjectFuture = SocketAccess.doPrivileged(() -> { InputStream inputStream = null; CompletableFuture putObjectRespFuture; try { @@ -657,6 +445,7 @@ private void uploadInOneChunk( } InputStream finalInputStream = inputStream; + return putObjectRespFuture.handle((resp, throwable) -> { releaseResourcesSafely(semaphore, finalInputStream, uploadRequest.getKey()); @@ -676,7 +465,7 @@ private void uploadInOneChunk( } catch (IOException e) { throw new RuntimeException(e); } - returnFuture.complete(null); + returnFuture.complete(ConditionalWriteResponse.success(resp.eTag())); } return null; @@ -694,106 +483,33 @@ private void uploadInOneChunk( CompletableFutureUtils.forwardResultTo(putObjectFuture, returnFuture); } - private void uploadInOneChunkConditionally( - S3AsyncClient s3AsyncClient, - UploadRequest uploadRequest, - StreamContext streamContext, + /** + exception handling for both regular and conditional uploads + */ + private void handleUploadException( CompletableFuture returnFuture, - StatsMetricPublisher statsMetricPublisher, - Semaphore semaphore + Throwable throwable, + String resourceName, + boolean isConditional ) { - PutObjectRequest.Builder putObjectRequestBuilder = PutObjectRequest.builder() - .bucket(uploadRequest.getBucket()) - .key(uploadRequest.getKey()) - .contentLength(uploadRequest.getContentLength()) - .overrideConfiguration(o -> o.addMetricPublisher(statsMetricPublisher.putObjectMetricPublisher)); - - if (CollectionUtils.isNotEmpty(uploadRequest.getMetadata())) { - putObjectRequestBuilder.metadata(uploadRequest.getMetadata()); - } - if (uploadRequest.doRemoteDataIntegrityCheck()) { - putObjectRequestBuilder.checksumAlgorithm(ChecksumAlgorithm.CRC32); - putObjectRequestBuilder.checksumCRC32(base64StringFromLong(uploadRequest.getExpectedChecksum())); - } - - if (uploadRequest.getConditionalOptions() != null) { - applyConditionalHeaders(putObjectRequestBuilder, uploadRequest.getConditionalOptions()); - } - - PutObjectRequest putObjectRequest = putObjectRequestBuilder.build(); - ExecutorService streamReadExecutor; - if (uploadRequest.getWritePriority() == WritePriority.URGENT) { - streamReadExecutor = urgentExecutorService; - } else if (uploadRequest.getWritePriority() == WritePriority.HIGH) { - streamReadExecutor = priorityExecutorService; + if (isConditional) { + handleConditionalException(returnFuture, throwable, resourceName); } else { - streamReadExecutor = executorService; - } - - CompletableFuture putObjectFuture = SocketAccess.doPrivileged(() -> { - InputStream inputStream = null; - CompletableFuture putObjectRespFuture; - try { - InputStreamContainer inputStreamContainer = streamContext.provideStream(0); - inputStream = AsyncPartsHandler.maybeRetryInputStream( - inputStreamContainer.getInputStream(), - uploadRequest.getWritePriority(), - uploadRequest.isUploadRetryEnabled(), - uploadRequest.getContentLength(), - maxRetryablePartSize - ); - AsyncRequestBody asyncRequestBody = AsyncRequestBody.fromInputStream( - inputStream, - inputStreamContainer.getContentLength(), - streamReadExecutor - ); - putObjectRespFuture = s3AsyncClient.putObject(putObjectRequest, asyncRequestBody); - } catch (Exception e) { - releaseResourcesSafely(semaphore, inputStream, uploadRequest.getKey()); - return CompletableFuture.failedFuture(e); - } - - InputStream finalInputStream = inputStream; - - return putObjectRespFuture.handle((resp, throwable) -> { - releaseResourcesSafely(semaphore, finalInputStream, uploadRequest.getKey()); - - if (throwable != null) { - Throwable unwrappedThrowable = ExceptionsHelper.unwrap(throwable, S3Exception.class); - if (unwrappedThrowable != null) { - S3Exception s3Exception = (S3Exception) unwrappedThrowable; - if (s3Exception.statusCode() == HttpStatusCode.BAD_REQUEST - && "BadDigest".equals(s3Exception.awsErrorDetails().errorCode())) { - throw new RuntimeException(new CorruptFileException(s3Exception, uploadRequest.getKey())); - } - } - returnFuture.completeExceptionally(throwable); + CompletableFuture voidFuture = new CompletableFuture<>(); + handleException(voidFuture, () -> "Upload failed for " + resourceName, throwable); + voidFuture.whenComplete((result, ex) -> { + if (ex != null) { + returnFuture.completeExceptionally(ex); } else { - try { - uploadRequest.getUploadFinalizer().accept(true); - } catch (IOException e) { - throw new RuntimeException(e); - } - returnFuture.complete(ConditionalWriteResponse.success(resp.eTag())); - } - - return null; - }).handle((resp, throwable) -> { - if (throwable != null) { - deleteUploadedObject(s3AsyncClient, uploadRequest); - returnFuture.completeExceptionally(throwable); + returnFuture.complete(ConditionalWriteResponse.success("")); } - - return null; }); - }); - - CompletableFutureUtils.forwardExceptionTo(returnFuture, putObjectFuture); - CompletableFutureUtils.forwardResultTo(putObjectFuture, returnFuture); + } } /** - * Apply conditional headers to a request builder + + Apply conditional headers to a request builder */ private void applyConditionalHeaders(Object builder, ConditionalWriteOptions options) { if (options == null) { @@ -817,6 +533,107 @@ private void applyConditionalHeaders(Object builder, ConditionalWriteOptions opt } } + /** + + Error handler for conditional uploads + */ + private void handleConditionalException(CompletableFuture returnFuture, Throwable throwable, String resourceName) { + Throwable cause = throwable; + while (cause instanceof CompletionException && cause.getCause() != null) { + cause = cause.getCause(); + } + + if (cause instanceof S3Exception s3e) { + if (s3e.statusCode() == HTTP_STATUS_PRECONDITION_FAILED) { + returnFuture.completeExceptionally( + S3Exception.builder() + .message("Conditional write failed: condition not met for " + resourceName) + .statusCode(HTTP_STATUS_PRECONDITION_FAILED) + .cause(s3e) + .build() + ); + return; + } else if (s3e.statusCode() == HTTP_STATUS_CONFLICT) { + returnFuture.completeExceptionally( + S3Exception.builder() + .message("Blob already exists: " + resourceName) + .statusCode(HTTP_STATUS_CONFLICT) + .cause(s3e) + .build() + ); + return; + } + } + if (cause instanceof Error) { + returnFuture.completeExceptionally(cause); + } else { + SdkClientException exception = SdkClientException.create("Failed conditional upload of " + resourceName, cause); + returnFuture.completeExceptionally(exception); + } + } + + private void mergeAndVerifyChecksum( + AtomicReferenceArray inputStreamContainers, + String fileName, + long expectedChecksum + ) { + long resultantChecksum = fromBase64String(inputStreamContainers.get(0).getChecksum()); + for (int index = 1; index < inputStreamContainers.length(); index++) { + long curChecksum = fromBase64String(inputStreamContainers.get(index).getChecksum()); + resultantChecksum = JZlib.crc32_combine(resultantChecksum, curChecksum, inputStreamContainers.get(index).getContentLength()); + } + + if (resultantChecksum != expectedChecksum) { + throw new RuntimeException(new CorruptFileException("File level checksums didn't match combined part checksums", fileName)); + } + + } + + private static String base64StringFromLong(Long val) { + return Base64.getEncoder().encodeToString(Arrays.copyOfRange(ByteUtils.toByteArrayBE(val), 4, 8)); + } + + private static long fromBase64String(String base64String) { + byte[] decodedBytes = Base64.getDecoder().decode(base64String); + if (decodedBytes.length != 4) { + throw new IllegalArgumentException("Invalid Base64 encoded CRC32 checksum"); + } + long result = 0; + for (int i = 0; i < 4; i++) { + result <<= 8; + result |= (decodedBytes[i] & 0xFF); + } + return result; + } + + private static void handleException(CompletableFuture returnFuture, Supplier message, Throwable throwable) { + Throwable cause = throwable instanceof CompletionException ? throwable.getCause() : throwable; + + if (cause instanceof Error) { + returnFuture.completeExceptionally(cause); + } else { + SdkClientException exception = SdkClientException.create(message.get(), cause); + returnFuture.completeExceptionally(exception); + } + + } + + /** + + Calculates the optimal part size of each part request if the upload operation is carried out as multipart upload. + */ + public long calculateOptimalPartSize(long contentLengthOfSource, WritePriority writePriority, boolean uploadRetryEnabled) { + if (contentLengthOfSource < ByteSizeUnit.MB.toBytes(5)) { + return contentLengthOfSource; + } + if (uploadRetryEnabled && (writePriority == WritePriority.HIGH || writePriority == WritePriority.URGENT)) { + return new ByteSizeValue(5, ByteSizeUnit.MB).getBytes(); + } + double optimalPartSize = contentLengthOfSource / (double) MAX_UPLOAD_PARTS; + optimalPartSize = Math.ceil(optimalPartSize); + return (long) Math.max(optimalPartSize, minimumPartSize); + } + private void releaseResourcesSafely(Semaphore semaphore, InputStream inputStream, String file) { if (semaphore != null) { semaphore.release(); From 656a3a1c522d023b60deeedd5d7400c0c9e6ed2b Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Wed, 11 Jun 2025 17:02:04 +0530 Subject: [PATCH 19/20] removed stale_primary_shard : generic rewording --- .../opensearch/repositories/s3/S3BlobStoreContainerTests.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java index 56c24fa43007d..ec2a05e33476c 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java @@ -1713,7 +1713,7 @@ public void testExecuteMultipartUploadConditionallyPreconditionFailed() { Exception exception = capturedException.get(); assertNotNull("Expected an exception to be captured", exception); assertTrue("Exception should be an OpenSearchException", exception instanceof OpenSearchException); - assertEquals("stale_primary_shard", exception.getMessage()); + assertEquals("Precondition Failed : Etag Mismatch", exception.getMessage()); verify(client).createMultipartUpload(any(CreateMultipartUploadRequest.class)); verify(client).completeMultipartUpload(any(CompleteMultipartUploadRequest.class)); @@ -1839,7 +1839,7 @@ public void testExecuteMultipartUploadConditionallyS3ExceptionTypes() { if ("S3Exception".equals(exceptionType) && statusCode == 412) { assertTrue(listenerException instanceof OpenSearchException); - assertEquals("stale_primary_shard", ((OpenSearchException) listenerException).getMessage()); + assertEquals("Precondition Failed : Etag Mismatch", listenerException.getMessage()); } else { assertTrue(listenerException instanceof IOException); } From 11dc26466ec5c5185d5bb9fca4107f93456a6cc6 Mon Sep 17 00:00:00 2001 From: Tanishq Ranjan Date: Thu, 12 Jun 2025 16:00:46 +0530 Subject: [PATCH 20/20] AsyncTransferManagerTests.java refactor --- .../repositories/s3/async/AsyncTransferManager.java | 6 ++---- .../s3/async/AsyncTransferManagerTests.java | 10 +++++----- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java index 94bbeeb90ceca..e7b8f57b474dd 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java @@ -516,15 +516,13 @@ private void applyConditionalHeaders(Object builder, ConditionalWriteOptions opt return; } - if (builder instanceof PutObjectRequest.Builder) { - PutObjectRequest.Builder putBuilder = (PutObjectRequest.Builder) builder; + if (builder instanceof PutObjectRequest.Builder putBuilder) { if (options.isIfNotExists()) { putBuilder.ifNoneMatch("*"); } else if (options.isIfMatch()) { putBuilder.ifMatch(options.getVersionIdentifier()); } - } else if (builder instanceof CompleteMultipartUploadRequest.Builder) { - CompleteMultipartUploadRequest.Builder completeBuilder = (CompleteMultipartUploadRequest.Builder) builder; + } else if (builder instanceof CompleteMultipartUploadRequest.Builder completeBuilder) { if (options.isIfNotExists()) { completeBuilder.ifNoneMatch("*"); } else if (options.isIfMatch()) { diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java index 9928211c54479..91383784a3b7f 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java @@ -338,7 +338,7 @@ public void testConditionalOneChunkUpload() { true, metadata, options, - null, + ServerSideEncryption.UNKNOWN_TO_SDK_VERSION.toString(), null, false, null, @@ -396,7 +396,7 @@ public void testConditionalOneChunkUploadPreconditionFailed() { true, null, options, - null, + ServerSideEncryption.UNKNOWN_TO_SDK_VERSION.toString(), null, false, null, @@ -466,7 +466,7 @@ public void testConditionalOneChunkUploadCorruption() { true, null, options, - null, + ServerSideEncryption.UNKNOWN_TO_SDK_VERSION.toString(), null, false, null, @@ -543,7 +543,7 @@ public void testConditionalMultipartUploadPreconditionFailed() { true, null, options, - null, + ServerSideEncryption.UNKNOWN_TO_SDK_VERSION.toString(), null, false, null, @@ -619,7 +619,7 @@ public void testConditionalMultipartUploadCorruption() { true, null, options, - null, + ServerSideEncryption.UNKNOWN_TO_SDK_VERSION.toString(), null, false, null,