Splitting a shard in-place

vikasvb90 · Jan 19, 2025 · f4db400 · f4db400
1 parent 34ef146
commit f4db400
Show file tree

Hide file tree

Showing 192 changed files with 8,619 additions and 808 deletions.
diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml
diff --git a/.idea/runConfigurations/Debug_OpenSearch.xml b/.idea/runConfigurations/Debug_OpenSearch.xml
diff --git a/...hmarks/src/main/java/org/opensearch/benchmark/routing/allocation/AllocationBenchmark.java b/...hmarks/src/main/java/org/opensearch/benchmark/routing/allocation/AllocationBenchmark.java
@@ -210,6 +210,7 @@ public void setUp() throws Exception {
         assert (initialClusterState.getRoutingNodes().shardsWithState(ShardRoutingState.STARTED).size() == totalShardCount);
         assert (initialClusterState.getRoutingNodes().shardsWithState(ShardRoutingState.INITIALIZING).size() == 0);
         assert (initialClusterState.getRoutingNodes().shardsWithState(ShardRoutingState.RELOCATING).size() == 0);
+        assert (initialClusterState.getRoutingNodes().shardsWithState(ShardRoutingState.SPLITTING).size() == 0);
         // make sure shards are only allocated on tag1
         for (ShardRouting startedShard : initialClusterState.getRoutingNodes().shardsWithState(ShardRoutingState.STARTED)) {
             assert (initialClusterState.getRoutingNodes().node(startedShard.currentNodeId()).node().getAttributes().get("tag")).equals(

diff --git a/...mmon/src/main/java/org/opensearch/common/annotation/processor/ApiAnnotationProcessor.java b/...mmon/src/main/java/org/opensearch/common/annotation/processor/ApiAnnotationProcessor.java
@@ -75,6 +75,9 @@ public Set<String> getSupportedOptions() {
 
     @Override
     public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment round) {
+        if (true) {
+            return false;
+        }
         processingEnv.getMessager().printMessage(Kind.NOTE, "Processing OpenSearch Api annotations");
 
         if (processingEnv.getOptions().containsKey(OPTION_CONTINUE_ON_FAILING_CHECKS) == true) {
@@ -167,6 +170,10 @@ private boolean validateVersion(String version) {
      * @param enclosing enclosing element
      */
     private void process(ExecutableElement executable, Element enclosing) {
+        if (true) {
+            return;
+        }
+
         if (!inspectable(executable)) {
             return;
         }
@@ -235,6 +242,9 @@ private void process(ExecutableElement executable, WildcardType type) {
      * @param ref reference type
      */
     private void process(ExecutableElement executable, ReferenceType ref) {
+        if (true) {
+            return;
+        }
         // The element has been processed already
         if (processed.add(ref) == false) {
             return;

diff --git a/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamInput.java b/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamInput.java
@@ -944,6 +944,18 @@ public long[] readVLongArray() throws IOException {
         return values;
     }
 
+    public long[] readZLongArray() throws IOException {
+        int length = readArraySize();
+        if (length == 0) {
+            return EMPTY_LONG_ARRAY;
+        }
+        long[] values = new long[length];
+        for (int i = 0; i < length; i++) {
+            values[i] = readZLong();
+        }
+        return values;
+    }
+
     private static final float[] EMPTY_FLOAT_ARRAY = new float[0];
 
     public float[] readFloatArray() throws IOException {

diff --git a/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamOutput.java b/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamOutput.java
@@ -907,6 +907,13 @@ public void writeVLongArray(long[] values) throws IOException {
         }
     }
 
+    public void writeZLongArray(long[] values) throws IOException {
+        writeVInt(values.length);
+        for (long value : values) {
+            writeZLong(value);
+        }
+    }
+
     public void writeFloatArray(float[] values) throws IOException {
         writeVInt(values.length);
         for (float value : values) {

diff --git a/plugins/repository-azure/src/main/java/org/opensearch/repositories/azure/AzureBlobStore.java b/plugins/repository-azure/src/main/java/org/opensearch/repositories/azure/AzureBlobStore.java
@@ -311,7 +311,7 @@ public Map<String, BlobMetadata> listBlobsByPrefix(String keyPath, String prefix
 
                 final BlobItemProperties properties = blobItem.getProperties();
                 logger.trace(() -> new ParameterizedMessage("blob name [{}], size [{}]", name, properties.getContentLength()));
-                blobsBuilder.put(name, new PlainBlobMetadata(name, properties.getContentLength()));
+                blobsBuilder.put(name, new PlainBlobMetadata(name, properties.getContentLength(), properties.getLastModified().toInstant().toEpochMilli()));
             }
         });
 

diff --git a/...sitory-gcs/src/main/java/org/opensearch/repositories/gcs/GoogleCloudStorageBlobStore.java b/...sitory-gcs/src/main/java/org/opensearch/repositories/gcs/GoogleCloudStorageBlobStore.java
@@ -172,7 +172,7 @@ Map<String, BlobMetadata> listBlobsByPrefix(String path, String prefix) throws I
                     assert blob.getName().startsWith(path);
                     if (blob.isDirectory() == false) {
                         final String suffixName = blob.getName().substring(path.length());
-                        mapBuilder.put(suffixName, new PlainBlobMetadata(suffixName, blob.getSize()));
+                        mapBuilder.put(suffixName, new PlainBlobMetadata(suffixName, blob.getSize(), blob.getUpdateTime()));
                     }
                 })
         );

diff --git a/...ins/repository-hdfs/src/main/java/org/opensearch/repositories/hdfs/HdfsBlobContainer.java b/...ins/repository-hdfs/src/main/java/org/opensearch/repositories/hdfs/HdfsBlobContainer.java
@@ -203,7 +203,7 @@ public Map<String, BlobMetadata> listBlobsByPrefix(@Nullable final String prefix
         Map<String, BlobMetadata> map = new LinkedHashMap<>();
         for (FileStatus file : files) {
             if (file.isFile()) {
-                map.put(file.getPath().getName(), new PlainBlobMetadata(file.getPath().getName(), file.getLen()));
+                map.put(file.getPath().getName(), new PlainBlobMetadata(file.getPath().getName(), file.getLen(), file.getModificationTime()));
             }
         }
         return Collections.unmodifiableMap(map);

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
@@ -32,6 +32,9 @@
 
 package org.opensearch.repositories.s3;
 
+import org.opensearch.cluster.metadata.CryptoMetadata;
+import org.opensearch.repositories.s3.async.CopyObjectHelper;
+import org.opensearch.repositories.s3.async.CopyObjectMetadata;
 import software.amazon.awssdk.core.ResponseInputStream;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.exception.SdkException;
@@ -53,6 +56,8 @@
 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.S3Error;
+import software.amazon.awssdk.services.s3.model.S3Object;
 import software.amazon.awssdk.services.s3.model.ServerSideEncryption;
 import software.amazon.awssdk.services.s3.model.UploadPartRequest;
 import software.amazon.awssdk.services.s3.model.UploadPartResponse;
@@ -96,11 +101,16 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
 import org.reactivestreams.Subscriber;
@@ -113,7 +123,7 @@
 class S3BlobContainer extends AbstractBlobContainer implements AsyncMultiStreamBlobContainer {
 
     private static final Logger logger = LogManager.getLogger(S3BlobContainer.class);
-
+    private static final long MAX_CONCURRENT_COPY_SIZE = ByteSizeUnit.GB.toBytes(20);
     private final S3BlobStore blobStore;
     private final String keyPath;
 
@@ -386,6 +396,75 @@ public void deleteBlobsIgnoringIfNotExists(List<String> blobNames) {
         future.actionGet();
     }
 
+    public Set<String> copyFilesFromSrcRemote(Set<String> files, AsyncMultiStreamBlobContainer sourceBlobContainer) {
+        if (!(sourceBlobContainer instanceof S3BlobContainer)) {
+            throw new IllegalArgumentException("Unknown blob container type. Not compatible with S3");
+        }
+        S3BlobContainer sourceS3Container = (S3BlobContainer) sourceBlobContainer;
+
+        Map<String, S3Object> fileObjects;
+        Set<String> filesFound = new HashSet<>(files.size());
+        try (AmazonS3Reference clientReference = sourceS3Container.blobStore.clientReference()) {
+            fileObjects = lisObjectsWithPredicate(clientReference, sourceS3Container, (file) -> {
+                if (files.contains(file)) {
+                    filesFound.add(file);
+                }
+                return files.contains(file);
+            });
+        }
+
+        try (AmazonAsyncS3Reference clientReference = sourceS3Container.blobStore.asyncClientReference()) {
+            CopyObjectHelper copyObjectHelper = new CopyObjectHelper(clientReference.get().client());
+            long curSize = 0;
+            List<CompletableFuture<Void>> copyAllFilesFutures = new ArrayList<>();
+            for (String remoteFileName : fileObjects.keySet()) {
+                S3Object fileObject = fileObjects.get(remoteFileName);
+                if (blobExists(fileObject.key()) == false) {
+                    curSize += fileObject.size();
+                    CopyObjectMetadata metadata = copyObjectHelper.fetchCopyObjectMetadata(
+                        sourceS3Container.blobStore.bucket(),
+                        fileObject.key());
+                    CompletableFuture<Void> copyFuture = copyObjectHelper.copyObject(
+                        sourceS3Container.blobStore.bucket(),
+                        fileObject.key(),
+                        metadata,
+                        blobStore.bucket(),
+                        buildKey(remoteFileName)
+                    );
+                    copyAllFilesFutures.add(copyFuture);
+
+                    if (curSize > MAX_CONCURRENT_COPY_SIZE) {
+                        waitForUploads(copyAllFilesFutures);
+                        copyAllFilesFutures.clear();
+                        curSize = 0;
+                    }
+                }
+            }
+
+            if (!copyAllFilesFutures.isEmpty()) {
+                waitForUploads(copyAllFilesFutures);
+            }
+        }
+        return filesFound;
+    }
+
+    @SuppressWarnings("rawtypes")
+    private static void waitForUploads(List<CompletableFuture<Void>> copyAllFilesFutures) {
+        CompletableFuture<Void> respFuture = CompletableFuture.allOf(copyAllFilesFutures.toArray(
+            new CompletableFuture[0])).thenApply(__ -> {
+            copyAllFilesFutures.forEach(CompletableFuture::join);
+            return null;
+        });
+
+        try {
+            respFuture.join();
+        } catch (Exception ex) {
+            logger.error("Failed to perform a remote copy of files", ex);
+            throw ex;
+        }
+    }
+
+
     @Override
     public List<BlobMetadata> listBlobsByPrefixInSortedOrder(String blobNamePrefix, int limit, BlobNameSortOrder blobNameSortOrder)
         throws IOException {
@@ -402,7 +481,7 @@ public List<BlobMetadata> listBlobsByPrefixInSortedOrder(String blobNamePrefix,
             try (AmazonS3Reference clientReference = blobStore.clientReference()) {
                 List<BlobMetadata> blobs = executeListing(clientReference, listObjectsRequest(prefix, limit), limit).stream()
                     .flatMap(listing -> listing.contents().stream())
-                    .map(s3Object -> new PlainBlobMetadata(s3Object.key().substring(keyPath.length()), s3Object.size()))
+                    .map(s3Object -> new PlainBlobMetadata(s3Object.key().substring(keyPath.length()), s3Object.size(), s3Object.lastModified().toEpochMilli()))
                     .collect(Collectors.toList());
                 return blobs.subList(0, Math.min(limit, blobs.size()));
             } catch (final Exception e) {
@@ -415,9 +494,9 @@ public List<BlobMetadata> listBlobsByPrefixInSortedOrder(String blobNamePrefix,
     public Map<String, BlobMetadata> listBlobsByPrefix(@Nullable String blobNamePrefix) throws IOException {
         String prefix = blobNamePrefix == null ? keyPath : buildKey(blobNamePrefix);
         try (AmazonS3Reference clientReference = blobStore.clientReference()) {
-            return executeListing(clientReference, listObjectsRequest(prefix)).stream()
+            return executeListing(clientReference, listObjectsRequest(prefix, blobStore)).stream()
                 .flatMap(listing -> listing.contents().stream())
-                .map(s3Object -> new PlainBlobMetadata(s3Object.key().substring(keyPath.length()), s3Object.size()))
+                .map(s3Object -> new PlainBlobMetadata(s3Object.key().substring(keyPath.length()), s3Object.size(), s3Object.lastModified().toEpochMilli()))
                 .collect(Collectors.toMap(PlainBlobMetadata::name, Function.identity()));
         } catch (final SdkException e) {
             throw new IOException("Exception when listing blobs by prefix [" + prefix + "]", e);
@@ -432,7 +511,7 @@ public Map<String, BlobMetadata> listBlobs() throws IOException {
     @Override
     public Map<String, BlobContainer> children() throws IOException {
         try (AmazonS3Reference clientReference = blobStore.clientReference()) {
-            return executeListing(clientReference, listObjectsRequest(keyPath)).stream().flatMap(listObjectsResponse -> {
+            return executeListing(clientReference, listObjectsRequest(keyPath, blobStore)).stream().flatMap(listObjectsResponse -> {
                 assert listObjectsResponse.contents().stream().noneMatch(s -> {
                     for (CommonPrefix commonPrefix : listObjectsResponse.commonPrefixes()) {
                         if (s.key().substring(keyPath.length()).startsWith(commonPrefix.prefix())) {
@@ -477,7 +556,29 @@ private static List<ListObjectsV2Response> executeListing(
         });
     }
 
-    private ListObjectsV2Request listObjectsRequest(String keyPath) {
+    private Map<String, S3Object> lisObjectsWithPredicate(
+        AmazonS3Reference clientReference,
+        S3BlobContainer blobContainer,
+        Predicate<String> includeFileObject
+    ) {
+        return SocketAccess.doPrivileged(() -> {
+            final Map<String, S3Object> results = new HashMap<>();
+            ListObjectsV2Request listObjectsRequest = listObjectsRequest(blobContainer.keyPath, blobContainer.blobStore);
+            ListObjectsV2Iterable listObjectsIterable = clientReference.get().listObjectsV2Paginator(listObjectsRequest);
+            for (ListObjectsV2Response listObjectsV2Response : listObjectsIterable) {
+                for (S3Object s3Object : listObjectsV2Response.contents()) {
+                    String remoteFileName = s3Object.key().substring(keyPath.length());
+                    if (includeFileObject.test(remoteFileName) == true) {
+                        results.put(remoteFileName, s3Object);
+                    }
+                }
+            }
+
+            return results;
+        });
+    }
+
+    private ListObjectsV2Request listObjectsRequest(String keyPath, S3BlobStore blobStore) {
         return ListObjectsV2Request.builder()
             .bucket(blobStore.bucket())
             .prefix(keyPath)
@@ -487,7 +588,7 @@ private ListObjectsV2Request listObjectsRequest(String keyPath) {
     }
 
     private ListObjectsV2Request listObjectsRequest(String keyPath, int limit) {
-        return listObjectsRequest(keyPath).toBuilder().maxKeys(Math.min(limit, 1000)).build();
+        return listObjectsRequest(keyPath, blobStore).toBuilder().maxKeys(Math.min(limit, 1000)).build();
     }
 
     private String buildKey(String blobName) {