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 2cb11541d924f..53371cd1529ce 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 @@ -48,6 +48,7 @@ import software.amazon.awssdk.services.s3.model.CreateMultipartUploadResponse; import software.amazon.awssdk.services.s3.model.DeleteObjectsRequest; import software.amazon.awssdk.services.s3.model.DeleteObjectsResponse; +import software.amazon.awssdk.services.s3.model.DeletedObject; import software.amazon.awssdk.services.s3.model.GetObjectAttributesParts; import software.amazon.awssdk.services.s3.model.GetObjectAttributesRequest; import software.amazon.awssdk.services.s3.model.GetObjectAttributesResponse; @@ -92,7 +93,6 @@ import java.util.Collections; import java.util.Comparator; import java.util.HashMap; -import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Map; @@ -102,6 +102,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @@ -286,9 +287,8 @@ public int numberOfPagesFetched() { } } - public void testDelete() throws IOException { + public void testDelete() throws Exception { final String bucketName = randomAlphaOfLengthBetween(1, 10); - final BlobPath blobPath = new BlobPath(); int bulkDeleteSize = 5; @@ -297,147 +297,314 @@ public void testDelete() throws IOException { when(blobStore.getStatsMetricPublisher()).thenReturn(new StatsMetricPublisher()); when(blobStore.getBulkDeletesSize()).thenReturn(bulkDeleteSize); - final S3Client client = mock(S3Client.class); - doAnswer(invocation -> new AmazonS3Reference(client)).when(blobStore).clientReference(); + final S3AsyncClient s3AsyncClient = mock(S3AsyncClient.class); + final AmazonAsyncS3Reference asyncClientReference = mock(AmazonAsyncS3Reference.class); + when(blobStore.asyncClientReference()).thenReturn(asyncClientReference); + AmazonAsyncS3WithCredentials amazonAsyncS3WithCredentials = AmazonAsyncS3WithCredentials.create( + s3AsyncClient, + s3AsyncClient, + s3AsyncClient, + null + ); + when(asyncClientReference.get()).thenReturn(amazonAsyncS3WithCredentials); - ListObjectsV2Iterable listObjectsV2Iterable = mock(ListObjectsV2Iterable.class); + final ListObjectsV2Publisher listPublisher = mock(ListObjectsV2Publisher.class); final int totalPageCount = 3; final long s3ObjectSize = ByteSizeUnit.MB.toBytes(5); final int s3ObjectsPerPage = 5; - MockListObjectsV2ResponseIterator listObjectsV2ResponseIterator = new MockListObjectsV2ResponseIterator( - totalPageCount, - s3ObjectsPerPage, - s3ObjectSize - ); - when(listObjectsV2Iterable.iterator()).thenReturn(listObjectsV2ResponseIterator); - when(client.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listObjectsV2Iterable); - final List keysDeleted = new ArrayList<>(); - AtomicInteger deleteCount = new AtomicInteger(); + List responses = new ArrayList<>(); + List allObjects = new ArrayList<>(); + long totalSize = 0; + + for (int i = 0; i < totalPageCount; i++) { + List pageObjects = new ArrayList<>(); + for (int j = 0; j < s3ObjectsPerPage; j++) { + pageObjects.add(S3Object.builder().key(randomAlphaOfLength(10)).size(s3ObjectSize).build()); + totalSize += s3ObjectSize; + } + allObjects.addAll(pageObjects); + responses.add(ListObjectsV2Response.builder().contents(pageObjects).build()); + } + + AtomicInteger counter = new AtomicInteger(); doAnswer(invocation -> { - DeleteObjectsRequest deleteObjectsRequest = invocation.getArgument(0); - deleteCount.getAndIncrement(); - logger.info("Object sizes are{}", deleteObjectsRequest.delete().objects().size()); - keysDeleted.addAll(deleteObjectsRequest.delete().objects().stream().map(ObjectIdentifier::key).collect(Collectors.toList())); - return DeleteObjectsResponse.builder().build(); - }).when(client).deleteObjects(any(DeleteObjectsRequest.class)); + Subscriber subscriber = invocation.getArgument(0); + subscriber.onSubscribe(new Subscription() { + @Override + public void request(long n) { + int currentCounter = counter.getAndIncrement(); + if (currentCounter < responses.size()) { + subscriber.onNext(responses.get(currentCounter)); + } + if (currentCounter == responses.size() - 1) { + subscriber.onComplete(); + } + } + + @Override + public void cancel() {} + }); + return null; + }).when(listPublisher).subscribe(ArgumentMatchers.>any()); + + when(s3AsyncClient.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listPublisher); + + when(s3AsyncClient.deleteObjects(any(DeleteObjectsRequest.class))).thenReturn( + CompletableFuture.completedFuture(DeleteObjectsResponse.builder().build()) + ); final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); - DeleteResult deleteResult = blobContainer.delete(); - assertEquals(s3ObjectSize * s3ObjectsPerPage * totalPageCount, deleteResult.bytesDeleted()); - assertEquals(s3ObjectsPerPage * totalPageCount, deleteResult.blobsDeleted()); - // keysDeleted will have blobPath also - assertEquals(listObjectsV2ResponseIterator.getKeysListed().size(), keysDeleted.size() - 1); - assertTrue(keysDeleted.contains(blobPath.buildAsString())); - // keysDeleted will have blobPath also - assertEquals((int) Math.ceil(((double) keysDeleted.size() + 1) / bulkDeleteSize), deleteCount.get()); - keysDeleted.remove(blobPath.buildAsString()); - assertEquals(new HashSet<>(listObjectsV2ResponseIterator.getKeysListed()), new HashSet<>(keysDeleted)); + CountDownLatch latch = new CountDownLatch(1); + AtomicReference resultRef = new AtomicReference<>(); + + blobContainer.deleteAsync(new ActionListener<>() { + @Override + public void onResponse(DeleteResult deleteResult) { + resultRef.set(deleteResult); + latch.countDown(); + } + + @Override + public void onFailure(Exception e) { + fail("Unexpected failure: " + e.getMessage()); + } + }); + + assertTrue(latch.await(5, TimeUnit.SECONDS)); + DeleteResult result = resultRef.get(); + + assertEquals(totalSize, result.bytesDeleted()); + assertEquals(allObjects.size(), result.blobsDeleted()); + + verify(s3AsyncClient, times(1)).listObjectsV2Paginator(any(ListObjectsV2Request.class)); + int expectedDeleteCalls = (int) Math.ceil((double) allObjects.size() / bulkDeleteSize); + verify(s3AsyncClient, times(expectedDeleteCalls)).deleteObjects(any(DeleteObjectsRequest.class)); } - public void testDeleteItemLevelErrorsDuringDelete() { + public void testDeleteItemLevelErrorsDuringDelete() throws Exception { final String bucketName = randomAlphaOfLengthBetween(1, 10); - final BlobPath blobPath = new BlobPath(); final S3BlobStore blobStore = mock(S3BlobStore.class); when(blobStore.bucket()).thenReturn(bucketName); when(blobStore.getStatsMetricPublisher()).thenReturn(new StatsMetricPublisher()); + int bulkDeleteSize = 3; // Small size to force multiple delete requests + when(blobStore.getBulkDeletesSize()).thenReturn(bulkDeleteSize); - final S3Client client = mock(S3Client.class); - doAnswer(invocation -> new AmazonS3Reference(client)).when(blobStore).clientReference(); + final S3AsyncClient s3AsyncClient = mock(S3AsyncClient.class); + final AmazonAsyncS3Reference asyncClientReference = mock(AmazonAsyncS3Reference.class); + when(blobStore.asyncClientReference()).thenReturn(asyncClientReference); + when(asyncClientReference.get()).thenReturn(AmazonAsyncS3WithCredentials.create(s3AsyncClient, s3AsyncClient, s3AsyncClient, null)); - ListObjectsV2Iterable listObjectsV2Iterable = mock(ListObjectsV2Iterable.class); - final int totalPageCount = 3; - final long s3ObjectSize = ByteSizeUnit.MB.toBytes(5); - final int s3ObjectsPerPage = 5; - MockListObjectsV2ResponseIterator listObjectsV2ResponseIterator = new MockListObjectsV2ResponseIterator( - totalPageCount, - s3ObjectsPerPage, - s3ObjectSize - ); - when(listObjectsV2Iterable.iterator()).thenReturn(listObjectsV2ResponseIterator); - when(client.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listObjectsV2Iterable); + final ListObjectsV2Publisher listPublisher = mock(ListObjectsV2Publisher.class); + final int totalObjects = 10; + List s3Objects = new ArrayList<>(); + for (int i = 0; i < totalObjects; i++) { + s3Objects.add(S3Object.builder().key("key-" + i).size(100L).build()); + } - final List keysFailedDeletion = new ArrayList<>(); + AtomicBoolean onNext = new AtomicBoolean(false); doAnswer(invocation -> { - DeleteObjectsRequest deleteObjectsRequest = invocation.getArgument(0); - int i = 0; - for (ObjectIdentifier objectIdentifier : deleteObjectsRequest.delete().objects()) { + Subscriber subscriber = invocation.getArgument(0); + subscriber.onSubscribe(new Subscription() { + @Override + public void request(long n) { + if (onNext.compareAndSet(false, true)) { + subscriber.onNext(ListObjectsV2Response.builder().contents(s3Objects).build()); + } else { + subscriber.onComplete(); + } + } + + @Override + public void cancel() {} + }); + return null; + }).when(listPublisher).subscribe(ArgumentMatchers.>any()); + + when(s3AsyncClient.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listPublisher); + + // Simulate item-level errors during delete + AtomicInteger deleteCallCount = new AtomicInteger(0); + when(s3AsyncClient.deleteObjects(any(DeleteObjectsRequest.class))).thenAnswer(invocation -> { + DeleteObjectsRequest request = invocation.getArgument(0); + List errors = new ArrayList<>(); + List deletedObjects = new ArrayList<>(); + + for (int i = 0; i < request.delete().objects().size(); i++) { if (i % 2 == 0) { - keysFailedDeletion.add(objectIdentifier.key()); + errors.add( + S3Error.builder() + .key(request.delete().objects().get(i).key()) + .code("InternalError") + .message("Simulated error") + .build() + ); + } else { + deletedObjects.add(DeletedObject.builder().key(request.delete().objects().get(i).key()).build()); } - i++; } - return DeleteObjectsResponse.builder() - .errors(keysFailedDeletion.stream().map(key -> S3Error.builder().key(key).build()).collect(Collectors.toList())) - .build(); - }).when(client).deleteObjects(any(DeleteObjectsRequest.class)); + + deleteCallCount.incrementAndGet(); + return CompletableFuture.completedFuture(DeleteObjectsResponse.builder().errors(errors).deleted(deletedObjects).build()); + }); final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); - assertThrows(AssertionError.class, blobContainer::delete); + CountDownLatch latch = new CountDownLatch(1); + AtomicReference resultRef = new AtomicReference<>(); + AtomicReference exceptionRef = new AtomicReference<>(); + + blobContainer.deleteAsync(new ActionListener<>() { + @Override + public void onResponse(DeleteResult deleteResult) { + resultRef.set(deleteResult); + latch.countDown(); + } + + @Override + public void onFailure(Exception e) { + exceptionRef.set(e); + latch.countDown(); + } + }); + + assertTrue(latch.await(5, TimeUnit.SECONDS)); + + assertNull("Unexpected exception: " + exceptionRef.get(), exceptionRef.get()); + DeleteResult result = resultRef.get(); + assertNotNull("Expected DeleteResult but got null", result); + + // We expect half of the objects to be deleted successfully + // But as of today, the blob delete count and bytes is updated a bit earlier. + assertEquals(totalObjects, result.blobsDeleted()); + assertEquals(totalObjects * 100L, result.bytesDeleted()); + + verify(s3AsyncClient, times(1)).listObjectsV2Paginator(any(ListObjectsV2Request.class)); + + // Calculate expected number of deleteObjects calls + int expectedDeleteCalls = (int) Math.ceil((double) totalObjects / bulkDeleteSize); + assertEquals(expectedDeleteCalls, deleteCallCount.get()); } - public void testDeleteSdkExceptionDuringListOperation() { + public void testDeleteSdkExceptionDuringListOperation() throws Exception { final String bucketName = randomAlphaOfLengthBetween(1, 10); - final BlobPath blobPath = new BlobPath(); final S3BlobStore blobStore = mock(S3BlobStore.class); when(blobStore.bucket()).thenReturn(bucketName); when(blobStore.getStatsMetricPublisher()).thenReturn(new StatsMetricPublisher()); - final S3Client client = mock(S3Client.class); - doAnswer(invocation -> new AmazonS3Reference(client)).when(blobStore).clientReference(); + final S3AsyncClient s3AsyncClient = mock(S3AsyncClient.class); + final AmazonAsyncS3Reference asyncClientReference = mock(AmazonAsyncS3Reference.class); + when(blobStore.asyncClientReference()).thenReturn(asyncClientReference); + when(asyncClientReference.get()).thenReturn(AmazonAsyncS3WithCredentials.create(s3AsyncClient, s3AsyncClient, s3AsyncClient, null)); - ListObjectsV2Iterable listObjectsV2Iterable = mock(ListObjectsV2Iterable.class); - final int totalPageCount = 3; - final long s3ObjectSize = ByteSizeUnit.MB.toBytes(5); - final int s3ObjectsPerPage = 5; - MockListObjectsV2ResponseIterator listObjectsV2ResponseIterator = new MockListObjectsV2ResponseIterator( - totalPageCount, - s3ObjectsPerPage, - s3ObjectSize - ); - when(listObjectsV2Iterable.iterator()).thenReturn(listObjectsV2ResponseIterator); - when(client.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listObjectsV2Iterable); + final ListObjectsV2Publisher listPublisher = mock(ListObjectsV2Publisher.class); + doAnswer(invocation -> { + Subscriber subscriber = invocation.getArgument(0); + subscriber.onSubscribe(new Subscription() { + @Override + public void request(long n) { + subscriber.onError(new RuntimeException("Simulated listing error")); + } + + @Override + public void cancel() {} + }); + return null; + }).when(listPublisher).subscribe(ArgumentMatchers.>any()); + + when(s3AsyncClient.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listPublisher); final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); - assertThrows(IOException.class, blobContainer::delete); + CountDownLatch latch = new CountDownLatch(1); + AtomicReference exceptionRef = new AtomicReference<>(); + + blobContainer.deleteAsync(new ActionListener<>() { + @Override + public void onResponse(DeleteResult deleteResult) { + fail("Expected failure but got success"); + } + + @Override + public void onFailure(Exception e) { + exceptionRef.set(e); + latch.countDown(); + } + }); + + assertTrue(latch.await(5, TimeUnit.SECONDS)); + assertNotNull(exceptionRef.get()); + assertEquals(IOException.class, exceptionRef.get().getClass()); + assertEquals("Failed to list objects for deletion", exceptionRef.get().getMessage()); } - public void testDeleteSdkExceptionDuringDeleteOperation() { + public void testDeleteSdkExceptionDuringDeleteOperation() throws Exception { final String bucketName = randomAlphaOfLengthBetween(1, 10); - final BlobPath blobPath = new BlobPath(); + int bulkDeleteSize = 5; final S3BlobStore blobStore = mock(S3BlobStore.class); when(blobStore.bucket()).thenReturn(bucketName); + when(blobStore.getBulkDeletesSize()).thenReturn(bulkDeleteSize); when(blobStore.getStatsMetricPublisher()).thenReturn(new StatsMetricPublisher()); - final S3Client client = mock(S3Client.class); - doAnswer(invocation -> new AmazonS3Reference(client)).when(blobStore).clientReference(); + final S3AsyncClient s3AsyncClient = mock(S3AsyncClient.class); + final AmazonAsyncS3Reference asyncClientReference = mock(AmazonAsyncS3Reference.class); + when(blobStore.asyncClientReference()).thenReturn(asyncClientReference); + when(asyncClientReference.get()).thenReturn(AmazonAsyncS3WithCredentials.create(s3AsyncClient, s3AsyncClient, s3AsyncClient, null)); - ListObjectsV2Iterable listObjectsV2Iterable = mock(ListObjectsV2Iterable.class); - final int totalPageCount = 3; - final long s3ObjectSize = ByteSizeUnit.MB.toBytes(5); - final int s3ObjectsPerPage = 5; - MockListObjectsV2ResponseIterator listObjectsV2ResponseIterator = new MockListObjectsV2ResponseIterator( - totalPageCount, - s3ObjectsPerPage, - s3ObjectSize - ); - when(listObjectsV2Iterable.iterator()).thenReturn(listObjectsV2ResponseIterator); - when(client.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listObjectsV2Iterable); + final ListObjectsV2Publisher listPublisher = mock(ListObjectsV2Publisher.class); + doAnswer(invocation -> { + Subscriber subscriber = invocation.getArgument(0); + subscriber.onSubscribe(new Subscription() { + @Override + public void request(long n) { + subscriber.onNext( + ListObjectsV2Response.builder().contents(S3Object.builder().key("test-key").size(100L).build()).build() + ); + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + return null; + }).when(listPublisher).subscribe(ArgumentMatchers.>any()); - when(client.deleteObjects(any(DeleteObjectsRequest.class))).thenThrow(SdkException.builder().build()); + when(s3AsyncClient.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listPublisher); + + CompletableFuture failedFuture = new CompletableFuture<>(); + failedFuture.completeExceptionally(new RuntimeException("Simulated delete error")); + when(s3AsyncClient.deleteObjects(any(DeleteObjectsRequest.class))).thenReturn(failedFuture); final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); - assertThrows(IOException.class, blobContainer::delete); + CountDownLatch latch = new CountDownLatch(1); + AtomicReference exceptionRef = new AtomicReference<>(); + + blobContainer.deleteAsync(new ActionListener<>() { + @Override + public void onResponse(DeleteResult deleteResult) { + fail("Expected failure but got success"); + } + + @Override + public void onFailure(Exception e) { + exceptionRef.set(e); + latch.countDown(); + } + }); + + assertTrue(latch.await(5, TimeUnit.SECONDS)); + assertNotNull(exceptionRef.get()); + logger.error("", exceptionRef.get()); + assertTrue(exceptionRef.get() instanceof CompletionException); + assertEquals("java.lang.RuntimeException: Simulated delete error", exceptionRef.get().getMessage()); } public void testExecuteSingleUpload() throws IOException {