Exercise multipart copy randomly in testWriteMaybeCopyRead

Author: bcully

modules/repository-s3/src/internalClusterTest/java/org/elasticsearch/repositories/s3/S3BlobStoreRepositoryTests.java

@@ -622,6 +622,12 @@ long getLargeBlobThresholdInBytes() {
                                     return ByteSizeUnit.MB.toBytes(1L);
                                 }
 
+                                @Override
+                                long getLargeCopyThresholdInBytes() {
+                                    // on my laptop 10K exercises this better but larger values should be fine for nightlies
+                                    return ByteSizeUnit.MB.toBytes(1L);
+                                }
+
                                 @Override
                                 void ensureMultiPartUploadSize(long blobSize) {}
                             };

modules/repository-s3/src/main/java/org/elasticsearch/repositories/s3/S3BlobContainer.java

@@ -304,6 +304,11 @@ long getLargeBlobThresholdInBytes() {
         return blobStore.bufferSizeInBytes();
     }
 
+    // package private for testing
+    long getLargeCopyThresholdInBytes() {
+        return MAX_FILE_SIZE.getBytes();
+    }
+
     @Override
     public void writeBlobAtomic(
         OperationPurpose purpose,
@@ -353,7 +358,7 @@ public void copyBlob(
         final var s3SourceBlobContainer = (S3BlobContainer) sourceBlobContainer;
 
         try {
-            if (blobSize > MAX_FILE_SIZE.getBytes()) {
+            if (blobSize > getLargeCopyThresholdInBytes()) {
                 executeMultipartCopy(purpose, s3SourceBlobContainer, sourceBlobName, blobName, blobSize);
             } else {
                 // metadata is inherited from source, but not canned ACL or storage class

test/fixtures/s3-fixture/src/main/java/fixture/s3/S3HttpHandler.java

@@ -155,10 +155,33 @@ public void handle(final HttpExchange exchange) throws IOException {
                 if (upload == null) {
                     exchange.sendResponseHeaders(RestStatus.NOT_FOUND.getStatus(), -1);
                 } else {
-                    final Tuple<String, BytesReference> blob = parseRequestBody(exchange);
-                    upload.addPart(blob.v1(), blob.v2());
-                    exchange.getResponseHeaders().add("ETag", blob.v1());
-                    exchange.sendResponseHeaders(RestStatus.OK.getStatus(), -1);
+                    // CopyPart is UploadPart with an x-amz-copy-source header
+                    final var sourceBlobName = exchange.getRequestHeaders().get("X-amz-copy-source");
+                    if (sourceBlobName != null) {
+                        var sourceBlob = blobs.get(sourceBlobName.getFirst());
+                        if (sourceBlob == null) {
+                            exchange.sendResponseHeaders(RestStatus.NOT_FOUND.getStatus(), -1);
+                        } else {
+                            var range = parsePartRange(exchange);
+                            // we'll assume for tests that the source object on the heap is under 2G
+                            var part = sourceBlob.slice(range.v1().intValue(), range.v2().intValue());
+                            var etag = UUIDs.randomBase64UUID();
+                            upload.addPart(etag, part);
+                            byte[] response = ("""
+                                <?xml version="1.0" encoding="UTF-8"?>
+                                <CopyPartResult>
+                                    <ETag>%s</ETag>
+                                </CopyPartResult>""".formatted(etag)).getBytes(StandardCharsets.UTF_8);
+                            exchange.getResponseHeaders().add("Content-Type", "application/xml");
+                            exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
+                            exchange.getResponseBody().write(response);
+                        }
+                    } else {
+                        final Tuple<String, BytesReference> blob = parseRequestBody(exchange);
+                        upload.addPart(blob.v1(), blob.v2());
+                        exchange.getResponseHeaders().add("ETag", blob.v1());
+                        exchange.sendResponseHeaders(RestStatus.OK.getStatus(), -1);
+                    }
                 }
 
             } else if (request.isCompleteMultipartUploadRequest()) {
@@ -205,14 +228,18 @@ public void handle(final HttpExchange exchange) throws IOException {
                 final var sourceBlobName = exchange.getRequestHeaders().get("X-amz-copy-source");
                 if (sourceBlobName != null) {
                     var sourceBlob = blobs.get(sourceBlobName.getFirst());
-                    blobs.put(request.path(), sourceBlob);
+                    if (sourceBlob == null) {
+                        exchange.sendResponseHeaders(RestStatus.NOT_FOUND.getStatus(), -1);
+                    } else {
+                        blobs.put(request.path(), sourceBlob);
 
-                    byte[] response = ("""
-                        <?xml version="1.0" encoding="UTF-8"?>
-                        <CopyObjectResult></CopyObjectResult>""").getBytes(StandardCharsets.UTF_8);
-                    exchange.getResponseHeaders().add("Content-Type", "application/xml");
-                    exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
-                    exchange.getResponseBody().write(response);
+                        byte[] response = ("""
+                            <?xml version="1.0" encoding="UTF-8"?>
+                            <CopyObjectResult></CopyObjectResult>""").getBytes(StandardCharsets.UTF_8);
+                        exchange.getResponseHeaders().add("Content-Type", "application/xml");
+                        exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
+                        exchange.getResponseBody().write(response);
+                    }
                 } else {
                     final Tuple<String, BytesReference> blob = parseRequestBody(exchange);
                     blobs.put(request.path(), blob.v2());
@@ -481,6 +508,27 @@ static List<String> extractPartEtags(BytesReference completeMultipartUploadBody)
         }
     }
 
+    private static final Pattern rangePattern = Pattern.compile("^bytes=([0-9]+)-([0-9]+)$");
+
+    private static Tuple<Long, Long> parsePartRange(final HttpExchange exchange) {
+        final var sourceRangeHeaders = exchange.getRequestHeaders().get("X-amz-copy-source-range");
+        if (sourceRangeHeaders == null) {
+            throw new IllegalStateException("missing x-amz-copy-source-range header");
+        }
+        if (sourceRangeHeaders.size() != 1) {
+            throw new IllegalStateException("expected 1 x-amz-copy-source-range header, found " + sourceRangeHeaders.size());
+        }
+        final var sourceRangeHeader = sourceRangeHeaders.getFirst();
+        final Matcher matcher = rangePattern.matcher(sourceRangeHeader);
+        if (matcher.find() == false) {
+            throw new IllegalStateException("invalid x-amz-copy-source-range header [" + sourceRangeHeader + "]");
+        }
+        final var start = Long.parseLong(matcher.group(1));
+        final var end = Long.parseLong(matcher.group(2));
+
+        return new Tuple<>(start, end - start + 1);
+    }
+
     MultipartUpload getUpload(String uploadId) {
         return uploads.get(uploadId);
     }