From 2629151e444e08f82000fe19f8cd6f165b6a22b7 Mon Sep 17 00:00:00 2001 From: Aditi Goyal Date: Wed, 25 Jun 2025 16:55:31 +0530 Subject: [PATCH 1/6] Changes for the Upload Workflow for merged segments --- .../index/shard/RemoteStoreUploader.java | 22 +++ .../shard/RemoteStoreUploaderService.java | 185 ++++++++++++++++++ 2 files changed, 207 insertions(+) create mode 100644 server/src/main/java/org/opensearch/index/shard/RemoteStoreUploader.java create mode 100644 server/src/main/java/org/opensearch/index/shard/RemoteStoreUploaderService.java diff --git a/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploader.java b/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploader.java new file mode 100644 index 0000000000000..48d3d38d00736 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploader.java @@ -0,0 +1,22 @@ +/* + * 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.index.shard; + +import org.opensearch.core.action.ActionListener; + +import java.util.Collection; +import java.util.Map; + +/** + * Interface to handle the functionality for upload data in the remote store + */ +public interface RemoteStoreUploader { + + void syncAndUploadNewSegments(Collection localSegments, Map localSegmentsSizeMap, ActionListener listener); +} diff --git a/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploaderService.java b/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploaderService.java new file mode 100644 index 0000000000000..28901b414dc3e --- /dev/null +++ b/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploaderService.java @@ -0,0 +1,185 @@ +/* + * 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.index.shard; + +import org.apache.logging.log4j.Logger; +import org.apache.logging.log4j.message.ParameterizedMessage; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FilterDirectory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.opensearch.action.support.GroupedActionListener; +import org.opensearch.cluster.routing.RecoverySource; +import org.opensearch.common.logging.Loggers; +import org.opensearch.common.util.UploadListener; +import org.opensearch.core.action.ActionListener; +import org.opensearch.index.remote.RemoteSegmentTransferTracker; +import org.opensearch.index.store.CompositeDirectory; +import org.opensearch.index.store.RemoteSegmentStoreDirectory; + +import java.io.IOException; +import java.util.Collection; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; + +/** + * The service essentially acts as a bridge between local segment storage and remote storage, + * ensuring efficient and reliable segment synchronization while providing comprehensive monitoring and error handling. + */ +public class RemoteStoreUploaderService implements RemoteStoreUploader { + + private final Logger logger; + + public static final Set EXCLUDE_FILES = Set.of("write.lock"); + + private final IndexShard indexShard; + private final Directory storeDirectory; + private final RemoteSegmentStoreDirectory remoteDirectory; + private final Map localSegmentChecksumMap; + // todo: check if we need to create a separate segment tracker as it is said to be linked to RL + private final RemoteSegmentTransferTracker segmentTracker; + + public RemoteStoreUploaderService( + IndexShard indexShard, + Directory storeDirectory, + RemoteSegmentStoreDirectory remoteDirectory, + RemoteSegmentTransferTracker segmentTracker + ) { + this.indexShard = indexShard; + // todo: check the prefix to be add for this class + logger = Loggers.getLogger(getClass(), indexShard.shardId()); + this.storeDirectory = storeDirectory; + this.remoteDirectory = remoteDirectory; + this.segmentTracker = segmentTracker; + this.localSegmentChecksumMap = new HashMap<>(); + } + + @Override + public void syncAndUploadNewSegments( + Collection localSegments, + Map localSegmentsSizeMap, + ActionListener listener + ) { + + Collection filteredFiles = localSegments.stream().filter(file -> !skipUpload(file)).toList(); + if (filteredFiles.isEmpty()) { + logger.debug("No new segments to upload in uploadNewSegments"); + listener.onResponse(null); + return; + } + + logger.debug("Effective new segments files to upload {}", filteredFiles); + ActionListener> mappedListener = ActionListener.map(listener, resp -> null); + GroupedActionListener batchUploadListener = new GroupedActionListener<>(mappedListener, filteredFiles.size()); + Directory directory = ((FilterDirectory) (((FilterDirectory) storeDirectory).getDelegate())).getDelegate(); + + for (String filteredFile : filteredFiles) { + // Initializing listener here to ensure that the stats increment operations are thread-safe + UploadListener statsListener = createUploadListener(localSegmentsSizeMap); + ActionListener aggregatedListener = ActionListener.wrap(resp -> { + statsListener.onSuccess(filteredFile); + batchUploadListener.onResponse(resp); + if (directory instanceof CompositeDirectory) { + ((CompositeDirectory) directory).afterSyncToRemote(filteredFile); + } + }, ex -> { + logger.warn(() -> new ParameterizedMessage("Exception: [{}] while uploading segment files", ex), ex); + if (ex instanceof CorruptIndexException) { + indexShard.failShard(ex.getMessage(), ex); + } + statsListener.onFailure(filteredFile); + batchUploadListener.onFailure(ex); + }); + statsListener.beforeUpload(filteredFile); + // Place where the actual upload is happening + remoteDirectory.copyFrom(storeDirectory, filteredFile, IOContext.DEFAULT, aggregatedListener, isLowPriorityUpload()); + } + } + + boolean isLowPriorityUpload() { + return isLocalOrSnapshotRecoveryOrSeeding(); + } + + boolean isLocalOrSnapshotRecoveryOrSeeding() { + // In this case when the primary mode is false, we need to upload segments to Remote Store + // This is required in case of remote migration seeding/snapshots/shrink/ split/clone where we need to durable persist + // all segments to remote before completing the recovery to ensure durability. + return (indexShard.state() == IndexShardState.RECOVERING && indexShard.shardRouting.primary()) + && indexShard.recoveryState() != null + && (indexShard.recoveryState().getRecoverySource().getType() == RecoverySource.Type.LOCAL_SHARDS + || indexShard.recoveryState().getRecoverySource().getType() == RecoverySource.Type.SNAPSHOT + || indexShard.shouldSeedRemoteStore()); + } + + /** + * Creates an {@link UploadListener} containing the stats population logic which would be triggered before and after segment upload events + * + * @param fileSizeMap updated map of current snapshot of local segments to their sizes + */ + private UploadListener createUploadListener(Map fileSizeMap) { + return new UploadListener() { + private long uploadStartTime = 0; + + @Override + public void beforeUpload(String file) { + // Start tracking the upload bytes started + segmentTracker.addUploadBytesStarted(fileSizeMap.get(file)); + uploadStartTime = System.currentTimeMillis(); + } + + @Override + public void onSuccess(String file) { + // Track upload success + segmentTracker.addUploadBytesSucceeded(fileSizeMap.get(file)); + segmentTracker.addToLatestUploadedFiles(file); + segmentTracker.addUploadTimeInMillis(Math.max(1, System.currentTimeMillis() - uploadStartTime)); + } + + @Override + public void onFailure(String file) { + // Track upload failure + segmentTracker.addUploadBytesFailed(fileSizeMap.get(file)); + segmentTracker.addUploadTimeInMillis(Math.max(1, System.currentTimeMillis() - uploadStartTime)); + } + }; + } + + /** + * Whether to upload a file or not depending on whether file is in excluded list or has been already uploaded. + * + * @param file that needs to be uploaded. + * @return true if the upload has to be skipped for the file. + */ + private boolean skipUpload(String file) { + try { + // Exclude files that are already uploaded and the exclude files to come up with the list of files to be uploaded. + // todo: Check if we need the second condition or is it just fail safe + return EXCLUDE_FILES.contains(file) || remoteDirectory.containsFile(file, getChecksumOfLocalFile(file)); + } catch (IOException e) { + logger.error( + "Exception while reading checksum of local segment file: {}, ignoring the exception and re-uploading the file", + file + ); + } + return false; + } + + private String getChecksumOfLocalFile(String file) throws IOException { + if (!localSegmentChecksumMap.containsKey(file)) { + try (IndexInput indexInput = storeDirectory.openInput(file, IOContext.READONCE)) { + String checksum = Long.toString(CodecUtil.retrieveChecksum(indexInput)); + localSegmentChecksumMap.put(file, checksum); + } + } + return localSegmentChecksumMap.get(file); + } +} From a9d8b2ac9e9461e2d6df775796ab1c88b13b2f81 Mon Sep 17 00:00:00 2001 From: Aditi Goyal Date: Thu, 26 Jun 2025 14:57:20 +0530 Subject: [PATCH 2/6] Addressed Todos --- .../index/shard/RemoteStoreUploader.java | 9 +- .../shard/RemoteStoreUploaderService.java | 115 +++--------------- 2 files changed, 25 insertions(+), 99 deletions(-) diff --git a/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploader.java b/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploader.java index 48d3d38d00736..39ce83215f43e 100644 --- a/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploader.java +++ b/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploader.java @@ -8,15 +8,22 @@ package org.opensearch.index.shard; +import org.opensearch.common.util.UploadListener; import org.opensearch.core.action.ActionListener; import java.util.Collection; import java.util.Map; +import java.util.function.Function; /** * Interface to handle the functionality for upload data in the remote store */ public interface RemoteStoreUploader { - void syncAndUploadNewSegments(Collection localSegments, Map localSegmentsSizeMap, ActionListener listener); + void uploadSegments( + Collection localSegments, + Map localSegmentsSizeMap, + ActionListener listener, + Function, UploadListener> uploadListenerFunction + ); } diff --git a/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploaderService.java b/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploaderService.java index 28901b414dc3e..b0e06ca978222 100644 --- a/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploaderService.java +++ b/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploaderService.java @@ -10,26 +10,21 @@ import org.apache.logging.log4j.Logger; import org.apache.logging.log4j.message.ParameterizedMessage; -import org.apache.lucene.codecs.CodecUtil; import org.apache.lucene.index.CorruptIndexException; import org.apache.lucene.store.Directory; import org.apache.lucene.store.FilterDirectory; import org.apache.lucene.store.IOContext; -import org.apache.lucene.store.IndexInput; import org.opensearch.action.support.GroupedActionListener; import org.opensearch.cluster.routing.RecoverySource; import org.opensearch.common.logging.Loggers; import org.opensearch.common.util.UploadListener; import org.opensearch.core.action.ActionListener; -import org.opensearch.index.remote.RemoteSegmentTransferTracker; import org.opensearch.index.store.CompositeDirectory; import org.opensearch.index.store.RemoteSegmentStoreDirectory; -import java.io.IOException; import java.util.Collection; -import java.util.HashMap; import java.util.Map; -import java.util.Set; +import java.util.function.Function; /** * The service essentially acts as a bridge between local segment storage and remote storage, @@ -39,69 +34,56 @@ public class RemoteStoreUploaderService implements RemoteStoreUploader { private final Logger logger; - public static final Set EXCLUDE_FILES = Set.of("write.lock"); - private final IndexShard indexShard; private final Directory storeDirectory; private final RemoteSegmentStoreDirectory remoteDirectory; - private final Map localSegmentChecksumMap; - // todo: check if we need to create a separate segment tracker as it is said to be linked to RL - private final RemoteSegmentTransferTracker segmentTracker; - public RemoteStoreUploaderService( - IndexShard indexShard, - Directory storeDirectory, - RemoteSegmentStoreDirectory remoteDirectory, - RemoteSegmentTransferTracker segmentTracker - ) { - this.indexShard = indexShard; - // todo: check the prefix to be add for this class + public RemoteStoreUploaderService(IndexShard indexShard, Directory storeDirectory, RemoteSegmentStoreDirectory remoteDirectory) { logger = Loggers.getLogger(getClass(), indexShard.shardId()); + this.indexShard = indexShard; this.storeDirectory = storeDirectory; this.remoteDirectory = remoteDirectory; - this.segmentTracker = segmentTracker; - this.localSegmentChecksumMap = new HashMap<>(); } @Override - public void syncAndUploadNewSegments( + public void uploadSegments( Collection localSegments, Map localSegmentsSizeMap, - ActionListener listener + ActionListener listener, + Function, UploadListener> uploadListenerFunction ) { - - Collection filteredFiles = localSegments.stream().filter(file -> !skipUpload(file)).toList(); - if (filteredFiles.isEmpty()) { + if (localSegments.isEmpty()) { logger.debug("No new segments to upload in uploadNewSegments"); listener.onResponse(null); return; } - logger.debug("Effective new segments files to upload {}", filteredFiles); + logger.debug("Effective new segments files to upload {}", localSegments); ActionListener> mappedListener = ActionListener.map(listener, resp -> null); - GroupedActionListener batchUploadListener = new GroupedActionListener<>(mappedListener, filteredFiles.size()); + GroupedActionListener batchUploadListener = new GroupedActionListener<>(mappedListener, localSegments.size()); Directory directory = ((FilterDirectory) (((FilterDirectory) storeDirectory).getDelegate())).getDelegate(); - for (String filteredFile : filteredFiles) { + for (String localSegment : localSegments) { // Initializing listener here to ensure that the stats increment operations are thread-safe - UploadListener statsListener = createUploadListener(localSegmentsSizeMap); + UploadListener statsListener = uploadListenerFunction.apply(localSegmentsSizeMap); ActionListener aggregatedListener = ActionListener.wrap(resp -> { - statsListener.onSuccess(filteredFile); + statsListener.onSuccess(localSegment); batchUploadListener.onResponse(resp); + // Once uploaded to Remote, local files become eligible for eviction from FileCache if (directory instanceof CompositeDirectory) { - ((CompositeDirectory) directory).afterSyncToRemote(filteredFile); + ((CompositeDirectory) directory).afterSyncToRemote(localSegment); } }, ex -> { logger.warn(() -> new ParameterizedMessage("Exception: [{}] while uploading segment files", ex), ex); if (ex instanceof CorruptIndexException) { indexShard.failShard(ex.getMessage(), ex); } - statsListener.onFailure(filteredFile); + statsListener.onFailure(localSegment); batchUploadListener.onFailure(ex); }); - statsListener.beforeUpload(filteredFile); + statsListener.beforeUpload(localSegment); // Place where the actual upload is happening - remoteDirectory.copyFrom(storeDirectory, filteredFile, IOContext.DEFAULT, aggregatedListener, isLowPriorityUpload()); + remoteDirectory.copyFrom(storeDirectory, localSegment, IOContext.DEFAULT, aggregatedListener, isLowPriorityUpload()); } } @@ -119,67 +101,4 @@ boolean isLocalOrSnapshotRecoveryOrSeeding() { || indexShard.recoveryState().getRecoverySource().getType() == RecoverySource.Type.SNAPSHOT || indexShard.shouldSeedRemoteStore()); } - - /** - * Creates an {@link UploadListener} containing the stats population logic which would be triggered before and after segment upload events - * - * @param fileSizeMap updated map of current snapshot of local segments to their sizes - */ - private UploadListener createUploadListener(Map fileSizeMap) { - return new UploadListener() { - private long uploadStartTime = 0; - - @Override - public void beforeUpload(String file) { - // Start tracking the upload bytes started - segmentTracker.addUploadBytesStarted(fileSizeMap.get(file)); - uploadStartTime = System.currentTimeMillis(); - } - - @Override - public void onSuccess(String file) { - // Track upload success - segmentTracker.addUploadBytesSucceeded(fileSizeMap.get(file)); - segmentTracker.addToLatestUploadedFiles(file); - segmentTracker.addUploadTimeInMillis(Math.max(1, System.currentTimeMillis() - uploadStartTime)); - } - - @Override - public void onFailure(String file) { - // Track upload failure - segmentTracker.addUploadBytesFailed(fileSizeMap.get(file)); - segmentTracker.addUploadTimeInMillis(Math.max(1, System.currentTimeMillis() - uploadStartTime)); - } - }; - } - - /** - * Whether to upload a file or not depending on whether file is in excluded list or has been already uploaded. - * - * @param file that needs to be uploaded. - * @return true if the upload has to be skipped for the file. - */ - private boolean skipUpload(String file) { - try { - // Exclude files that are already uploaded and the exclude files to come up with the list of files to be uploaded. - // todo: Check if we need the second condition or is it just fail safe - return EXCLUDE_FILES.contains(file) || remoteDirectory.containsFile(file, getChecksumOfLocalFile(file)); - } catch (IOException e) { - logger.error( - "Exception while reading checksum of local segment file: {}, ignoring the exception and re-uploading the file", - file - ); - } - return false; - } - - private String getChecksumOfLocalFile(String file) throws IOException { - if (!localSegmentChecksumMap.containsKey(file)) { - try (IndexInput indexInput = storeDirectory.openInput(file, IOContext.READONCE)) { - String checksum = Long.toString(CodecUtil.retrieveChecksum(indexInput)); - localSegmentChecksumMap.put(file, checksum); - } - } - return localSegmentChecksumMap.get(file); - } } From 3be078d46172a30394ad40946aeb6fec5b4e47e4 Mon Sep 17 00:00:00 2001 From: kh3ra Date: Sat, 28 Jun 2025 20:28:24 +0530 Subject: [PATCH 3/6] changes --- .../opensearch/index/shard/IndexShard.java | 10 +- .../org/opensearch/indices/IndicesModule.java | 2 + .../ActiveMergesSegmentRegistry.java | 89 +++++++++ ...oint.java => MergedSegmentCheckpoint.java} | 8 +- .../checkpoint/MergedSegmentPublisher.java | 8 +- .../PublishMergedSegmentAction.java | 31 +-- .../PublishMergedSegmentActionProvider.java | 65 +++++++ .../PublishMergedSegmentRequest.java | 8 +- .../RemoteStoreMergedSegmentCheckpoint.java | 102 ++++++++++ ...RemoteStorePublishMergedSegmentAction.java | 184 ++++++++++++++++++ ...emoteStorePublishMergedSegmentRequest.java | 62 ++++++ .../MergedSegmentReplicationTargetTests.java | 6 +- .../PrimaryShardReplicationSourceTests.java | 6 +- .../PublishMergedSegmentActionTests.java | 8 +- .../PublishMergedSegmentRequestTests.java | 4 +- 15 files changed, 550 insertions(+), 43 deletions(-) create mode 100644 server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java rename server/src/main/java/org/opensearch/indices/replication/checkpoint/{MergeSegmentCheckpoint.java => MergedSegmentCheckpoint.java} (93%) create mode 100644 server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java create mode 100644 server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java create mode 100644 server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java create mode 100644 server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentRequest.java diff --git a/server/src/main/java/org/opensearch/index/shard/IndexShard.java b/server/src/main/java/org/opensearch/index/shard/IndexShard.java index 315360fa4d99e..9ce911d2a8ac6 100644 --- a/server/src/main/java/org/opensearch/index/shard/IndexShard.java +++ b/server/src/main/java/org/opensearch/index/shard/IndexShard.java @@ -201,7 +201,7 @@ import org.opensearch.indices.recovery.RecoverySettings; import org.opensearch.indices.recovery.RecoveryState; import org.opensearch.indices.recovery.RecoveryTarget; -import org.opensearch.indices.replication.checkpoint.MergeSegmentCheckpoint; +import org.opensearch.indices.replication.checkpoint.MergedSegmentCheckpoint; import org.opensearch.indices.replication.checkpoint.MergedSegmentPublisher; import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint; import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher; @@ -1864,19 +1864,19 @@ public void publishMergedSegment(SegmentCommitInfo segmentCommitInfo) throws IOE } /** - * Compute {@link MergeSegmentCheckpoint} from a SegmentCommitInfo. + * Compute {@link MergedSegmentCheckpoint} from a SegmentCommitInfo. * This function fetches a metadata snapshot from the store that comes with an IO cost. * * @param segmentCommitInfo {@link SegmentCommitInfo} segmentCommitInfo to use to compute. - * @return {@link MergeSegmentCheckpoint} Checkpoint computed from the segmentCommitInfo. + * @return {@link MergedSegmentCheckpoint} Checkpoint computed from the segmentCommitInfo. * @throws IOException When there is an error computing segment metadata from the store. */ - public MergeSegmentCheckpoint computeMergeSegmentCheckpoint(SegmentCommitInfo segmentCommitInfo) throws IOException { + public MergedSegmentCheckpoint computeMergeSegmentCheckpoint(SegmentCommitInfo segmentCommitInfo) throws IOException { // Only need to get the file metadata information in segmentCommitInfo and reuse Store#getSegmentMetadataMap. SegmentInfos segmentInfos = new SegmentInfos(Version.LATEST.major); segmentInfos.add(segmentCommitInfo); Map segmentMetadataMap = store.getSegmentMetadataMap(segmentInfos); - return new MergeSegmentCheckpoint( + return new MergedSegmentCheckpoint( shardId, getOperationPrimaryTerm(), segmentMetadataMap.values().stream().mapToLong(StoreFileMetadata::length).sum(), diff --git a/server/src/main/java/org/opensearch/indices/IndicesModule.java b/server/src/main/java/org/opensearch/indices/IndicesModule.java index ad2eb834b721e..117e56cb8dd8d 100644 --- a/server/src/main/java/org/opensearch/indices/IndicesModule.java +++ b/server/src/main/java/org/opensearch/indices/IndicesModule.java @@ -83,6 +83,7 @@ import org.opensearch.indices.cluster.IndicesClusterStateService; import org.opensearch.indices.mapper.MapperRegistry; import org.opensearch.indices.replication.checkpoint.MergedSegmentPublisher; +import org.opensearch.indices.replication.checkpoint.PublishMergedSegmentActionProvider; import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher; import org.opensearch.indices.store.IndicesStore; import org.opensearch.indices.store.TransportNodesListShardStoreMetadata; @@ -299,6 +300,7 @@ protected void configure() { bind(RetentionLeaseBackgroundSyncAction.class).asEagerSingleton(); bind(RetentionLeaseSyncer.class).asEagerSingleton(); bind(SegmentReplicationCheckpointPublisher.class).asEagerSingleton(); + bind(MergedSegmentPublisher.PublishAction.class).toProvider(PublishMergedSegmentActionProvider.class).asEagerSingleton(); bind(MergedSegmentPublisher.class).asEagerSingleton(); bind(SegmentReplicationPressureService.class).asEagerSingleton(); bind(RemoteStorePressureService.class).asEagerSingleton(); diff --git a/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java b/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java new file mode 100644 index 0000000000000..672a875f9ad00 --- /dev/null +++ b/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java @@ -0,0 +1,89 @@ +/* + * 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.indices.replication; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import reactor.util.annotation.NonNull; + +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; + +/** + * TODO@kheraadi: Simplify this + * Usages: + * 1. Used by GC to exclude merged segment files for deletion + * 2. Used to track remote store segment file names (remove this) + * + */ +public class ActiveMergesSegmentRegistry { + private final Map myMergedSegments = new ConcurrentHashMap<>(); + private static final String PENDING_REMOTE_FILE_NAME = "PENDING_REMOTE_FILE_NAME"; + + private static class HOLDER { + private static final ActiveMergesSegmentRegistry INSTANCE = new ActiveMergesSegmentRegistry(); + } + + private ActiveMergesSegmentRegistry() {}; + + public static ActiveMergesSegmentRegistry getInstance() { + return HOLDER.INSTANCE; + } + + public void updateRemoteSegmentFileName(@NonNull String localSegmentFileName, @NonNull String remoteSegmentFileName) { + if (contains(localSegmentFileName) && PENDING_REMOTE_FILE_NAME.equals(getExistingRemoteSegmentFileName(localSegmentFileName)) == false){ + // This should never happen + if (remoteSegmentFileName.equals(getExistingRemoteSegmentFileName(localSegmentFileName))) { + return; + } + throw new IllegalArgumentException("Segment " + localSegmentFileName + " is already registered as " + getExistingRemoteSegmentFileName(localSegmentFileName) + ". Called with " + remoteSegmentFileName); + } + myMergedSegments.put(localSegmentFileName, remoteSegmentFileName); + } + + public void register(@NonNull String localSegmentFileName) { + if (contains(localSegmentFileName)){ + if(PENDING_REMOTE_FILE_NAME.equals(getExistingRemoteSegmentFileName(localSegmentFileName))) { + return; + } + throw new IllegalArgumentException(localSegmentFileName + ": " + getExistingRemoteSegmentFileName(localSegmentFileName) + " already registered. Cannot reregister."); + } + + myMergedSegments.put(localSegmentFileName, PENDING_REMOTE_FILE_NAME); + } + + public void unregister(@NonNull String segmentFileName) { + myMergedSegments.remove(segmentFileName); + } + + public boolean contains(@NonNull String segmentFileName) { + return myMergedSegments.containsKey(segmentFileName); + } + + public String getExistingRemoteSegmentFileName(@NonNull String localSegmentFileName) { + if (contains(localSegmentFileName) == false) { + // This should never happen + throw new IllegalArgumentException("Segment " + localSegmentFileName + " is not registered"); + } + + return myMergedSegments.get(localSegmentFileName); + } + + public boolean canDelete(@NonNull String segmentFileName) { + String originalFileName = getOriginalFileName(segmentFileName); + return contains(originalFileName) && + segmentFileName.equals(getExistingRemoteSegmentFileName(originalFileName)); + } + + private String getOriginalFileName(@NonNull String remoteSegmentFileName) { + String originalFileName = remoteSegmentFileName.split("__")[0]; + return originalFileName; + } +} diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergeSegmentCheckpoint.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentCheckpoint.java similarity index 93% rename from server/src/main/java/org/opensearch/indices/replication/checkpoint/MergeSegmentCheckpoint.java rename to server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentCheckpoint.java index 424a2cc98939e..6b8c82cafea8a 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergeSegmentCheckpoint.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentCheckpoint.java @@ -50,10 +50,10 @@ * @opensearch.internal */ @ExperimentalApi -public class MergeSegmentCheckpoint extends ReplicationCheckpoint { +public class MergedSegmentCheckpoint extends ReplicationCheckpoint { private final String segmentName; - public MergeSegmentCheckpoint( + public MergedSegmentCheckpoint( ShardId shardId, long primaryTerm, long length, @@ -65,7 +65,7 @@ public MergeSegmentCheckpoint( this.segmentName = segmentName; } - public MergeSegmentCheckpoint(StreamInput in) throws IOException { + public MergedSegmentCheckpoint(StreamInput in) throws IOException { super(in); segmentName = in.readString(); } @@ -89,7 +89,7 @@ public void writeTo(StreamOutput out) throws IOException { public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; - MergeSegmentCheckpoint that = (MergeSegmentCheckpoint) o; + MergedSegmentCheckpoint that = (MergedSegmentCheckpoint) o; return getPrimaryTerm() == that.getPrimaryTerm() && segmentName.equals(that.segmentName) && Objects.equals(getShardId(), that.getShardId()) diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentPublisher.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentPublisher.java index 4b49e71093c55..4f58b10500d1e 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentPublisher.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentPublisher.java @@ -25,15 +25,11 @@ public class MergedSegmentPublisher { // This Component is behind feature flag so we are manually binding this in IndicesModule. @Inject - public MergedSegmentPublisher(PublishMergedSegmentAction publishAction) { - this(publishAction::publish); - } - public MergedSegmentPublisher(PublishAction publishAction) { this.publishAction = Objects.requireNonNull(publishAction); } - public void publish(IndexShard indexShard, MergeSegmentCheckpoint checkpoint) { + public void publish(IndexShard indexShard, MergedSegmentCheckpoint checkpoint) { publishAction.publish(indexShard, checkpoint); } @@ -44,7 +40,7 @@ public void publish(IndexShard indexShard, MergeSegmentCheckpoint checkpoint) { */ @ExperimentalApi public interface PublishAction { - void publish(IndexShard indexShard, MergeSegmentCheckpoint checkpoint); + void publish(IndexShard indexShard, ReplicationCheckpoint checkpoint); } /** diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentAction.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentAction.java index 675a13c5c05ae..d10b799ea8df0 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentAction.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentAction.java @@ -31,10 +31,11 @@ * @opensearch.api */ @ExperimentalApi -public class PublishMergedSegmentAction extends AbstractPublishCheckpointAction { +public class PublishMergedSegmentAction extends AbstractPublishCheckpointAction implements MergedSegmentPublisher.PublishAction { public static final String ACTION_NAME = "indices:admin/publish_merged_segment"; - protected static Logger logger = LogManager.getLogger(PublishMergedSegmentAction.class); + + private static final Logger logger = LogManager.getLogger(PublishMergedSegmentAction.class); private final SegmentReplicationTargetService replicationService; @@ -66,33 +67,39 @@ public PublishMergedSegmentAction( this.replicationService = targetService; } + + @Override protected void doExecute(Task task, PublishMergedSegmentRequest request, ActionListener listener) { assert false : "use PublishMergedSegmentAction#publish"; } + @Override + protected void shardOperationOnPrimary( + PublishMergedSegmentRequest request, + IndexShard primary, + ActionListener> listener + ) { + ActionListener.completeWith(listener, () -> new PrimaryResult<>(request, new ReplicationResponse())); + } + /** * Publish merged segment request to shard */ - final void publish(IndexShard indexShard, MergeSegmentCheckpoint checkpoint) { + @Override + public void publish(IndexShard indexShard, ReplicationCheckpoint checkpoint) { + assert checkpoint instanceof MergedSegmentCheckpoint; doPublish( indexShard, checkpoint, - new PublishMergedSegmentRequest(checkpoint), + new PublishMergedSegmentRequest((MergedSegmentCheckpoint) checkpoint), "segrep_publish_merged_segment", true, indexShard.getRecoverySettings().getMergedSegmentReplicationTimeout() ); } - @Override - protected void shardOperationOnPrimary( - PublishMergedSegmentRequest request, - IndexShard primary, - ActionListener> listener - ) { - ActionListener.completeWith(listener, () -> new PrimaryResult<>(request, new ReplicationResponse())); - } + @Override protected void doReplicaOperation(PublishMergedSegmentRequest request, IndexShard replica) { diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java new file mode 100644 index 0000000000000..49f033e0264d9 --- /dev/null +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java @@ -0,0 +1,65 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.indices.replication.checkpoint; + +import org.opensearch.action.support.ActionFilters; +import org.opensearch.cluster.action.shard.ShardStateAction; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.inject.Inject; +import org.opensearch.common.inject.Provider; +import org.opensearch.common.settings.Settings; +import org.opensearch.indices.IndicesService; +import org.opensearch.indices.replication.SegmentReplicationTargetService; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +public class PublishMergedSegmentActionProvider implements Provider { + + private final Settings settings; + private final TransportService transportService; + private final ClusterService clusterService; + private final IndicesService indicesService; + private final ThreadPool threadPool; + private final ShardStateAction shardStateAction; + private final ActionFilters actionFilters; + private final SegmentReplicationTargetService targetService; + + @Inject + public PublishMergedSegmentActionProvider( + Settings settings, + TransportService transportService, + ClusterService clusterService, + IndicesService indicesService, + ThreadPool threadPool, + ShardStateAction shardStateAction, + ActionFilters actionFilters, + SegmentReplicationTargetService targetService + ) { + this.settings = settings; + this.transportService = transportService; + this.clusterService = clusterService; + this.indicesService = indicesService; + this.threadPool = threadPool; + this.shardStateAction = shardStateAction; + this.actionFilters = actionFilters; + this.targetService = targetService; + } + + @Override + public MergedSegmentPublisher.PublishAction get() { + // TODO@kheraadi: check if remote store is enabled here + if (false) { + return new PublishMergedSegmentAction( + settings, transportService, clusterService, indicesService, + threadPool, shardStateAction, actionFilters, targetService + ); + } else { + return new RemoteStorePublishMergedSegmentAction( + settings, transportService, clusterService, indicesService, + threadPool, shardStateAction, actionFilters, targetService + ); + } + } +} diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentRequest.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentRequest.java index 0dbb54846522a..1ebc167134a3f 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentRequest.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentRequest.java @@ -21,16 +21,16 @@ * @opensearch.internal */ public class PublishMergedSegmentRequest extends ReplicationRequest { - private final MergeSegmentCheckpoint mergedSegment; + private final MergedSegmentCheckpoint mergedSegment; - public PublishMergedSegmentRequest(MergeSegmentCheckpoint mergedSegment) { + public PublishMergedSegmentRequest(MergedSegmentCheckpoint mergedSegment) { super(mergedSegment.getShardId()); this.mergedSegment = mergedSegment; } public PublishMergedSegmentRequest(StreamInput in) throws IOException { super(in); - this.mergedSegment = new MergeSegmentCheckpoint(in); + this.mergedSegment = new MergedSegmentCheckpoint(in); } @Override @@ -56,7 +56,7 @@ public String toString() { return "PublishMergedSegmentRequest{" + "mergedSegment=" + mergedSegment + '}'; } - public MergeSegmentCheckpoint getMergedSegment() { + public MergedSegmentCheckpoint getMergedSegment() { return mergedSegment; } } diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java new file mode 100644 index 0000000000000..605dcf18e94d6 --- /dev/null +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java @@ -0,0 +1,102 @@ +/* + * 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.indices.replication.checkpoint; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.seqno.SequenceNumbers; +import org.opensearch.index.store.StoreFileMetadata; + +import java.io.IOException; +import java.util.Map; +import java.util.Objects; + +/** + * Represents a remote store merged segment checkpoint. + * + * @opensearch.internal + */ +@ExperimentalApi +public class RemoteStoreMergedSegmentCheckpoint extends ReplicationCheckpoint { + private final String segmentName; + private final String remoteStorePath; + + public RemoteStoreMergedSegmentCheckpoint( + ShardId shardId, + long primaryTerm, + long length, + String codec, + Map metadataMap, + String segmentName, + String remoteStorePath + ) { + super(shardId, primaryTerm, SequenceNumbers.NO_OPS_PERFORMED, SequenceNumbers.NO_OPS_PERFORMED, length, codec, metadataMap); + this.segmentName = segmentName; + this.remoteStorePath = remoteStorePath; + } + + public RemoteStoreMergedSegmentCheckpoint(StreamInput in) throws IOException { + super(in); + segmentName = in.readString(); + remoteStorePath = in.readString(); + } + + public String getSegmentName() { + return segmentName; + } + + public String getRemoteStorePath() { + return remoteStorePath; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(segmentName); + out.writeString(remoteStorePath); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + RemoteStoreMergedSegmentCheckpoint that = (RemoteStoreMergedSegmentCheckpoint) o; + return getPrimaryTerm() == that.getPrimaryTerm() + && segmentName.equals(that.segmentName) + && remoteStorePath.equals(that.remoteStorePath) + && Objects.equals(getShardId(), that.getShardId()) + && getCodec().equals(that.getCodec()); + } + + @Override + public int hashCode() { + return Objects.hash(getShardId(), getPrimaryTerm(), segmentName, remoteStorePath); + } + + @Override + public String toString() { + return "RemoteStoreMergedSegmentCheckpoint{" + + "shardId=" + + getShardId() + + ", primaryTerm=" + + getPrimaryTerm() + + ", segmentName=" + + segmentName + + ", remoteStorePath=" + + remoteStorePath + + '}'; + } + + public Map getLocalToRemoteSegmentFileNameMap() { + return null; + } +} diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java new file mode 100644 index 0000000000000..9c2d64fd90008 --- /dev/null +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java @@ -0,0 +1,184 @@ +/* + * 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.indices.replication.checkpoint; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.replication.ReplicationResponse; +import org.opensearch.cluster.action.shard.ShardStateAction; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.inject.Inject; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.UploadListener; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.common.io.stream.Writeable; +import org.opensearch.index.shard.IndexShard; +import org.opensearch.indices.IndicesService; +import org.opensearch.indices.replication.ActiveMergesSegmentRegistry; +import org.opensearch.indices.replication.SegmentReplicationTargetService; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import java.util.Collection; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +public class RemoteStorePublishMergedSegmentAction extends AbstractPublishCheckpointAction implements MergedSegmentPublisher.PublishAction { + + public static final String ACTION_NAME = "indices:admin/publish_merged_segment"; + + private final static Logger logger = LogManager.getLogger(RemoteStorePublishMergedSegmentAction.class); + + private final ActiveMergesSegmentRegistry activeMergesSegmentRegistry = ActiveMergesSegmentRegistry.getInstance(); + + private final SegmentReplicationTargetService replicationService; + + @Inject + public RemoteStorePublishMergedSegmentAction( + Settings settings, + TransportService transportService, + ClusterService clusterService, + IndicesService indicesService, + ThreadPool threadPool, + ShardStateAction shardStateAction, + ActionFilters actionFilters, + SegmentReplicationTargetService targetService + ) { + super( + settings, + ACTION_NAME, + transportService, + clusterService, + indicesService, + threadPool, + shardStateAction, + actionFilters, + RemoteStorePublishMergedSegmentRequest::new, + RemoteStorePublishMergedSegmentRequest::new, + ThreadPool.Names.GENERIC, + logger + ); + this.replicationService = targetService; + } + + @Override + protected void doReplicaOperation(RemoteStorePublishMergedSegmentRequest shardRequest, IndexShard replica) { + + } + + @Override + protected void shardOperationOnPrimary(RemoteStorePublishMergedSegmentRequest shardRequest, IndexShard primary, ActionListener> listener) { + + } + + @Override + public void publish(IndexShard indexShard, ReplicationCheckpoint checkpoint) { + assert checkpoint instanceof RemoteStoreMergedSegmentCheckpoint; + RemoteStoreMergedSegmentCheckpoint mergedSegmentCheckpoint = (RemoteStoreMergedSegmentCheckpoint) checkpoint; + + publishMergedSegmentsToRemoteStore(indexShard, mergedSegmentCheckpoint); + logger.info("RemoteFileNames at {}#publish {}", getClass().getName(), mergedSegmentCheckpoint).getLocalToRemoteSegmentFileNameMap()); + doPublish(indexShard, + checkpoint, + new RemoteStorePublishMergedSegmentRequest((RemoteStoreMergedSegmentCheckpoint) checkpoint), + "segrep_publish_merged_segment", + true, + indexShard.getRecoverySettings().getMergedSegmentReplicationTimeout() + ); + mergedSegmentCheckpoint + .getLocalToRemoteSegmentFileNameMap() + .keySet() + .forEach(activeMergesSegmentRegistry::unregister); + } + + private void publishMergedSegmentsToRemoteStore(IndexShard indexShard, RemoteStoreMergedSegmentCheckpoint checkpoint) { + RemoteStoreUploaderService remoteStoreUploaderService = getRemoteStoreUploaderService(indexShard); + Collection segmentsToUpload = checkpoint.getMetadataMap().keySet(); + logger.info("Publishing segments {} to remote store", segmentsToUpload); + registerSegmentsToActiveMerges(segmentsToUpload); + Map segmentsSizeMap = checkpoint + .getMetadataMap() + .entrySet() + .stream() + .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().length())); + + final CountDownLatch latch = new CountDownLatch(segmentsToUpload.size()); + + remoteStoreUploaderService.uploadSegments( + segmentsToUpload, + segmentsSizeMap, + new ActionListener() { + @Override + public void onResponse(Void unused) { + logger.info("Successfully uploaded segments {} to remote store", checkpoint.getLocalToRemoteSegmentFileNameMap()); + } + + @Override + public void onFailure(Exception e) { + logger.error("Failed to upload segments {} to remote store", segmentsToUpload, e); + segmentsToUpload.forEach(activeMergesSegmentRegistry::unregister); + throw new RuntimeException(e); + } + }, + (x) -> new UploadListener() { + @Override + public void beforeUpload(String file) { + + } + + @Override + public void onSuccess(String file) { + logger.info("Uploaded {}", file); + checkpoint.addRemoteSegmentFileName( + file, + activeMergesSegmentRegistry.getExistingRemoteSegmentFileName(file) + ); + latch.countDown(); + } + + @Override + public void onFailure(String file) { + segmentsToUpload.forEach(activeMergesSegmentRegistry::unregister); + /** + * TODO@kheraadi: + * 1. reset ActiveMergesRegistry + * 2. abort merge + */ + } + } + ); + try { + if(latch.await(60, TimeUnit.MINUTES) == false) {throw new RuntimeException("Merged segment upload timed out.");}; // TODO@kheraadi: Finalize timeout + } catch (InterruptedException e) { + throw new RuntimeException(e); + // TODO@kheraadi: abort merge properly here + } + } + + /** + * TODO@kheraadi: REBASE ONCE UPLOAD CHANGES ARE COMPLETE + */ + private RemoteStoreUploaderService getRemoteStoreUploaderService(IndexShard indexShard) { + return new RemoteStoreUploaderService( + indexShard, + indexShard.store().directory(), + indexShard.getRemoteDirectory() + ); + } + + private void registerSegmentsToActiveMerges(Collection segmentsToUpload) { + logger.info("Registering segments to active merges: {}", segmentsToUpload); + segmentsToUpload.forEach(activeMergesSegmentRegistry::register); + logger.info("ActiveMergesSegmentsRegistry: " + activeMergesSegmentsRegistry.mergedSegments().toString()); + } + +} diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentRequest.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentRequest.java new file mode 100644 index 0000000000000..862855b740a2d --- /dev/null +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentRequest.java @@ -0,0 +1,62 @@ +/* + * 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.indices.replication.checkpoint; + +import org.opensearch.action.support.replication.ReplicationRequest; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; + +import java.io.IOException; +import java.util.Objects; + +/** + * Replication request responsible for publishing remote store merged segment request to a replica shard. + * + * @opensearch.internal + */ +public class RemoteStorePublishMergedSegmentRequest extends ReplicationRequest { + private final RemoteStoreMergedSegmentCheckpoint mergedSegment; + + public RemoteStorePublishMergedSegmentRequest(RemoteStoreMergedSegmentCheckpoint mergedSegment) { + super(mergedSegment.getShardId()); + this.mergedSegment = mergedSegment; + } + + public RemoteStorePublishMergedSegmentRequest(StreamInput in) throws IOException { + super(in); + this.mergedSegment = new RemoteStoreMergedSegmentCheckpoint(in); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + mergedSegment.writeTo(out); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (!(o instanceof RemoteStorePublishMergedSegmentRequest that)) return false; + return Objects.equals(mergedSegment, that.mergedSegment); + } + + @Override + public int hashCode() { + return Objects.hash(mergedSegment); + } + + @Override + public String toString() { + return "RemoteStorePublishMergedSegmentRequest{" + "mergedSegment=" + mergedSegment + '}'; + } + + public RemoteStoreMergedSegmentCheckpoint getMergedSegment() { + return mergedSegment; + } +} \ No newline at end of file diff --git a/server/src/test/java/org/opensearch/indices/replication/MergedSegmentReplicationTargetTests.java b/server/src/test/java/org/opensearch/indices/replication/MergedSegmentReplicationTargetTests.java index 53c8e29acf2c9..1d9e47a7fc2f4 100644 --- a/server/src/test/java/org/opensearch/indices/replication/MergedSegmentReplicationTargetTests.java +++ b/server/src/test/java/org/opensearch/indices/replication/MergedSegmentReplicationTargetTests.java @@ -22,7 +22,7 @@ import org.opensearch.index.shard.IndexShard; import org.opensearch.index.shard.IndexShardTestCase; import org.opensearch.index.store.StoreFileMetadata; -import org.opensearch.indices.replication.checkpoint.MergeSegmentCheckpoint; +import org.opensearch.indices.replication.checkpoint.MergedSegmentCheckpoint; import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint; import org.opensearch.indices.replication.common.ReplicationFailedException; import org.opensearch.indices.replication.common.ReplicationType; @@ -41,7 +41,7 @@ public class MergedSegmentReplicationTargetTests extends IndexShardTestCase { private MergedSegmentReplicationTarget mergedSegmentReplicationTarget; private IndexShard indexShard, spyIndexShard; - private MergeSegmentCheckpoint mergedSegment; + private MergedSegmentCheckpoint mergedSegment; private ByteBuffersDataOutput buffer; private static final String SEGMENT_NAME = "_0.si"; @@ -71,7 +71,7 @@ public void setUp() throws Exception { try (ByteBuffersIndexOutput indexOutput = new ByteBuffersIndexOutput(buffer, "", null)) { testSegmentInfos.write(indexOutput); } - mergedSegment = new MergeSegmentCheckpoint( + mergedSegment = new MergedSegmentCheckpoint( spyIndexShard.shardId(), spyIndexShard.getPendingPrimaryTerm(), 1, diff --git a/server/src/test/java/org/opensearch/indices/replication/PrimaryShardReplicationSourceTests.java b/server/src/test/java/org/opensearch/indices/replication/PrimaryShardReplicationSourceTests.java index f9b236fad5b02..69badc5389939 100644 --- a/server/src/test/java/org/opensearch/indices/replication/PrimaryShardReplicationSourceTests.java +++ b/server/src/test/java/org/opensearch/indices/replication/PrimaryShardReplicationSourceTests.java @@ -21,7 +21,7 @@ import org.opensearch.index.shard.IndexShardTestCase; import org.opensearch.index.store.StoreFileMetadata; import org.opensearch.indices.recovery.RecoverySettings; -import org.opensearch.indices.replication.checkpoint.MergeSegmentCheckpoint; +import org.opensearch.indices.replication.checkpoint.MergedSegmentCheckpoint; import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint; import org.opensearch.telemetry.tracing.noop.NoopTracer; import org.opensearch.test.ClusterServiceUtils; @@ -136,7 +136,7 @@ public void testGetSegmentFiles() { public void testGetMergedSegmentFiles() { StoreFileMetadata testMetadata = new StoreFileMetadata("testFile", 1L, "checksum", Version.LATEST); - final ReplicationCheckpoint checkpoint = new MergeSegmentCheckpoint( + final ReplicationCheckpoint checkpoint = new MergedSegmentCheckpoint( indexShard.shardId(), PRIMARY_TERM, 1, @@ -195,7 +195,7 @@ public void testTransportTimeoutForGetSegmentFilesAction() { public void testTransportTimeoutForGetMergedSegmentFilesAction() { long fileSize = (long) (Math.pow(10, 9)); StoreFileMetadata testMetadata = new StoreFileMetadata("testFile", fileSize, "checksum", Version.LATEST); - final ReplicationCheckpoint checkpoint = new MergeSegmentCheckpoint( + final ReplicationCheckpoint checkpoint = new MergedSegmentCheckpoint( indexShard.shardId(), PRIMARY_TERM, 1, diff --git a/server/src/test/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionTests.java b/server/src/test/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionTests.java index 38c2c3d54f95e..e983ed22926ef 100644 --- a/server/src/test/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionTests.java +++ b/server/src/test/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionTests.java @@ -129,7 +129,7 @@ public void testPublishMergedSegment() { mockTargetService ); - final MergeSegmentCheckpoint checkpoint = new MergeSegmentCheckpoint( + final MergedSegmentCheckpoint checkpoint = new MergedSegmentCheckpoint( indexShard.shardId(), 1, 1111, @@ -168,7 +168,7 @@ public void testPublishMergedSegmentActionOnPrimary() { mockTargetService ); - final MergeSegmentCheckpoint checkpoint = new MergeSegmentCheckpoint( + final MergedSegmentCheckpoint checkpoint = new MergedSegmentCheckpoint( indexShard.shardId(), 1, 1111, @@ -212,7 +212,7 @@ public void testPublishMergedSegmentActionOnReplica() { mockTargetService ); - final MergeSegmentCheckpoint checkpoint = new MergeSegmentCheckpoint( + final MergedSegmentCheckpoint checkpoint = new MergedSegmentCheckpoint( indexShard.shardId(), 1, 1111, @@ -263,7 +263,7 @@ public void testPublishMergedSegmentActionOnDocrepReplicaDuringMigration() { mockTargetService ); - final MergeSegmentCheckpoint checkpoint = new MergeSegmentCheckpoint( + final MergedSegmentCheckpoint checkpoint = new MergedSegmentCheckpoint( indexShard.shardId(), 1, 1111, diff --git a/server/src/test/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentRequestTests.java b/server/src/test/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentRequestTests.java index e039a598164b4..f089e8e566b4d 100644 --- a/server/src/test/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentRequestTests.java +++ b/server/src/test/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentRequestTests.java @@ -22,7 +22,7 @@ public class PublishMergedSegmentRequestTests extends OpenSearchTestCase { public void testPublishMergedSegmentRequest() { - MergeSegmentCheckpoint checkpoint = new MergeSegmentCheckpoint( + MergedSegmentCheckpoint checkpoint = new MergedSegmentCheckpoint( new ShardId(new Index("1", "1"), 0), 0, 0, @@ -37,7 +37,7 @@ public void testPublishMergedSegmentRequest() { } public void testSerialize() throws Exception { - MergeSegmentCheckpoint checkpoint = new MergeSegmentCheckpoint( + MergedSegmentCheckpoint checkpoint = new MergedSegmentCheckpoint( new ShardId(new Index("1", "1"), 0), 0, 0, From e68d228161b84256dc2451d9edf6fc64a61a0e37 Mon Sep 17 00:00:00 2001 From: kh3ra Date: Mon, 30 Jun 2025 17:09:26 +0530 Subject: [PATCH 4/6] Adding support for uploading merged segments in IndexWriter.warm --- .../engine/MergedSegmentWarmerFactory.java | 2 +- .../RemoteStoreMergedSegmentWarmer.java | 23 +++- .../opensearch/index/shard/IndexShard.java | 18 ++- .../shard/RemoteStoreRefreshListener.java | 4 + .../index/store/RemoteDirectory.java | 33 +++++- .../store/RemoteSegmentStoreDirectory.java | 65 ++++++++++- .../RemoteSegmentStoreDirectoryFactory.java | 3 +- .../store/RemoteStoreFileDownloader.java | 60 ++++++++-- .../ActiveMergesSegmentRegistry.java | 106 +++++++++++------- .../RemoteStoreReplicationSource.java | 55 +++++++++ .../SegmentReplicationTargetService.java | 2 +- .../AbstractPublishCheckpointAction.java | 8 +- .../checkpoint/MergedSegmentPublisher.java | 2 +- .../RemoteStoreMergedSegmentCheckpoint.java | 44 ++++---- ...RemoteStorePublishMergedSegmentAction.java | 41 +++---- .../repositories/FilterRepository.java | 5 + .../opensearch/repositories/Repository.java | 5 + .../blobstore/BlobStoreRepository.java | 24 ++++ 18 files changed, 383 insertions(+), 117 deletions(-) diff --git a/server/src/main/java/org/opensearch/index/engine/MergedSegmentWarmerFactory.java b/server/src/main/java/org/opensearch/index/engine/MergedSegmentWarmerFactory.java index eda93d22d9c3b..bb6493c955f99 100644 --- a/server/src/main/java/org/opensearch/index/engine/MergedSegmentWarmerFactory.java +++ b/server/src/main/java/org/opensearch/index/engine/MergedSegmentWarmerFactory.java @@ -59,7 +59,7 @@ public MergedSegmentWarmerFactory(TransportService transportService, RecoverySet public IndexWriter.IndexReaderWarmer get(IndexShard shard) { if (shard.indexSettings().isAssignedOnRemoteNode()) { - return new RemoteStoreMergedSegmentWarmer(transportService, recoverySettings, clusterService); + return new RemoteStoreMergedSegmentWarmer(transportService, recoverySettings, clusterService, shard); } else if (shard.indexSettings().isSegRepLocalEnabled()) { return new LocalMergedSegmentWarmer(transportService, recoverySettings, clusterService, shard); } else if (shard.indexSettings().isDocumentReplication()) { diff --git a/server/src/main/java/org/opensearch/index/engine/RemoteStoreMergedSegmentWarmer.java b/server/src/main/java/org/opensearch/index/engine/RemoteStoreMergedSegmentWarmer.java index f3d5f8a4cc93b..3853ab3fd3936 100644 --- a/server/src/main/java/org/opensearch/index/engine/RemoteStoreMergedSegmentWarmer.java +++ b/server/src/main/java/org/opensearch/index/engine/RemoteStoreMergedSegmentWarmer.java @@ -32,9 +32,14 @@ package org.opensearch.index.engine; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.SegmentCommitInfo; +import org.apache.lucene.index.SegmentReader; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.index.shard.IndexShard; import org.opensearch.indices.recovery.RecoverySettings; import org.opensearch.transport.TransportService; @@ -49,19 +54,33 @@ public class RemoteStoreMergedSegmentWarmer implements IndexWriter.IndexReaderWa private final TransportService transportService; private final RecoverySettings recoverySettings; private final ClusterService clusterService; + private final IndexShard indexShard; + + private final Logger logger = LogManager.getLogger(RemoteStoreMergedSegmentWarmer.class); public RemoteStoreMergedSegmentWarmer( TransportService transportService, RecoverySettings recoverySettings, - ClusterService clusterService + ClusterService clusterService, + IndexShard indexShard ) { this.transportService = transportService; this.recoverySettings = recoverySettings; this.clusterService = clusterService; + this.indexShard = indexShard; } @Override public void warm(LeafReader leafReader) throws IOException { - // TODO: remote store merged segment warmer + // IndexWriter.IndexReaderWarmer#warm is called by IndexWriter#mergeMiddle. The type of leafReader should be SegmentReader. + assert leafReader instanceof SegmentReader; + + SegmentCommitInfo segmentCommitInfo = ((SegmentReader) leafReader).getSegmentInfo(); + if(logger.isTraceEnabled()) { + logger.trace("[ShardId {}] Warming segment: {}", indexShard.shardId(), segmentCommitInfo); + } + indexShard.publishMergedSegment(segmentCommitInfo); + logger.trace("Completed segment warming for {} on shard {}", + segmentCommitInfo.info.name, indexShard.shardId()); } } diff --git a/server/src/main/java/org/opensearch/index/shard/IndexShard.java b/server/src/main/java/org/opensearch/index/shard/IndexShard.java index 9ce911d2a8ac6..c8a6f1d23610f 100644 --- a/server/src/main/java/org/opensearch/index/shard/IndexShard.java +++ b/server/src/main/java/org/opensearch/index/shard/IndexShard.java @@ -201,10 +201,7 @@ import org.opensearch.indices.recovery.RecoverySettings; import org.opensearch.indices.recovery.RecoveryState; import org.opensearch.indices.recovery.RecoveryTarget; -import org.opensearch.indices.replication.checkpoint.MergedSegmentCheckpoint; -import org.opensearch.indices.replication.checkpoint.MergedSegmentPublisher; -import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint; -import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher; +import org.opensearch.indices.replication.checkpoint.*; import org.opensearch.indices.replication.common.ReplicationTimer; import org.opensearch.repositories.RepositoriesService; import org.opensearch.repositories.Repository; @@ -1871,11 +1868,22 @@ public void publishMergedSegment(SegmentCommitInfo segmentCommitInfo) throws IOE * @return {@link MergedSegmentCheckpoint} Checkpoint computed from the segmentCommitInfo. * @throws IOException When there is an error computing segment metadata from the store. */ - public MergedSegmentCheckpoint computeMergeSegmentCheckpoint(SegmentCommitInfo segmentCommitInfo) throws IOException { + public ReplicationCheckpoint computeMergeSegmentCheckpoint(SegmentCommitInfo segmentCommitInfo) throws IOException { // Only need to get the file metadata information in segmentCommitInfo and reuse Store#getSegmentMetadataMap. SegmentInfos segmentInfos = new SegmentInfos(Version.LATEST.major); segmentInfos.add(segmentCommitInfo); Map segmentMetadataMap = store.getSegmentMetadataMap(segmentInfos); + if (indexSettings.isRemoteStoreEnabled()) { + return new RemoteStoreMergedSegmentCheckpoint( + shardId, + getOperationPrimaryTerm(), + segmentMetadataMap.values().stream().mapToLong(StoreFileMetadata::length).sum(), + getEngine().config().getCodec().getName(), + segmentMetadataMap, + segmentCommitInfo.info.name, + null + ); + } return new MergedSegmentCheckpoint( shardId, getOperationPrimaryTerm(), diff --git a/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java b/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java index 8ace4848806d7..ae02e20f9cc13 100644 --- a/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java +++ b/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java @@ -35,6 +35,7 @@ import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadata; import org.opensearch.index.translog.Translog; import org.opensearch.indices.RemoteStoreSettings; +import org.opensearch.indices.replication.ActiveMergesSegmentRegistry; import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint; import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher; import org.opensearch.threadpool.ThreadPool; @@ -84,6 +85,7 @@ public final class RemoteStoreRefreshListener extends ReleasableRetryableRefresh public static final Set EXCLUDE_FILES = Set.of("write.lock"); + private final ActiveMergesSegmentRegistry activeMergesSegmentRegistry = ActiveMergesSegmentRegistry.getInstance(); private final IndexShard indexShard; private final Directory storeDirectory; private final RemoteSegmentStoreDirectory remoteDirectory; @@ -258,6 +260,8 @@ private boolean syncSegments() { long lastRefreshedCheckpoint = ((InternalEngine) indexShard.getEngine()).lastRefreshedCheckpoint(); Collection localSegmentsPostRefresh = segmentInfos.files(true); + remoteDirectory.syncSegmentsUploadedToRemoteStoreWithActiveMergesSegmentRegistry(storeDirectory, localSegmentsPostRefresh); + // Create a map of file name to size and update the refresh segment tracker Map localSegmentsSizeMap = updateLocalSizeMapAndTracker(localSegmentsPostRefresh).entrySet() .stream() diff --git a/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java b/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java index 38d1b4d1252cc..45e1cf8193dca 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java @@ -32,6 +32,7 @@ import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.unit.ByteSizeUnit; import org.opensearch.index.store.exception.ChecksumCombinationException; +import org.opensearch.indices.replication.ActiveMergesSegmentRegistry; import java.io.FileNotFoundException; import java.io.IOException; @@ -67,7 +68,7 @@ public class RemoteDirectory extends Directory { private final UnaryOperator lowPriorityUploadRateLimiter; - private final UnaryOperator downloadRateLimiter; + private final DownloadRateLimiterProvider downloadRateLimiterProvider; /** * Number of bytes in the segment file to store checksum @@ -79,19 +80,20 @@ public BlobContainer getBlobContainer() { } public RemoteDirectory(BlobContainer blobContainer) { - this(blobContainer, UnaryOperator.identity(), UnaryOperator.identity(), UnaryOperator.identity()); + this(blobContainer, UnaryOperator.identity(), UnaryOperator.identity(), UnaryOperator.identity(), UnaryOperator.identity()); } public RemoteDirectory( BlobContainer blobContainer, UnaryOperator uploadRateLimiter, UnaryOperator lowPriorityUploadRateLimiter, - UnaryOperator downloadRateLimiter + UnaryOperator downloadRateLimiter, + UnaryOperator lowPriorityDownloadRateLimiter ) { this.blobContainer = blobContainer; this.lowPriorityUploadRateLimiter = lowPriorityUploadRateLimiter; this.uploadRateLimiter = uploadRateLimiter; - this.downloadRateLimiter = downloadRateLimiter; + this.downloadRateLimiterProvider = new DownloadRateLimiterProvider(downloadRateLimiter, lowPriorityDownloadRateLimiter); } /** @@ -236,7 +238,7 @@ public IndexInput openInput(String name, long fileLength, IOContext context) thr InputStream inputStream = null; try { inputStream = blobContainer.readBlob(name); - return new RemoteIndexInput(name, downloadRateLimiter.apply(inputStream), fileLength); + return new RemoteIndexInput(name, downloadRateLimiterProvider.get(name).apply(inputStream), fileLength); } catch (Exception e) { // In case the RemoteIndexInput creation fails, close the input stream to avoid file handler leak. if (inputStream != null) { @@ -475,8 +477,27 @@ private IndexInput getBlockInput(String name, long position, long length, long f byte[] bytes; try (InputStream inputStream = blobContainer.readBlob(name, position, length)) { // TODO - Explore how we can buffer small chunks of data instead of having the whole 8MB block in memory - bytes = downloadRateLimiter.apply(inputStream).readAllBytes(); + bytes = downloadRateLimiterProvider.get(name).apply(inputStream).readAllBytes(); } return new ByteArrayIndexInput(name, bytes); } + + private class DownloadRateLimiterProvider { + private final ActiveMergesSegmentRegistry activeMergesSegmentRegistry = ActiveMergesSegmentRegistry.getInstance(); + + private final UnaryOperator downloadRateLimiter; + private final UnaryOperator lowPriorityDownloadRateLimiter; + + DownloadRateLimiterProvider(UnaryOperator downloadRateLimiter, UnaryOperator lowPriorityDownloadRateLimiter) { + this.downloadRateLimiter = downloadRateLimiter; + this.lowPriorityDownloadRateLimiter = lowPriorityDownloadRateLimiter; + } + + public UnaryOperator get(final String filename){ + if (activeMergesSegmentRegistry.contains(filename)) { + return lowPriorityDownloadRateLimiter; + } + return downloadRateLimiter; + } + } } diff --git a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java index 1fd2e15ee50ec..5cc27f55d9892 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java @@ -39,6 +39,7 @@ import org.opensearch.index.store.lockmanager.RemoteStoreMetadataLockManager; import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadata; import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadataHandlerFactory; +import org.opensearch.indices.replication.ActiveMergesSegmentRegistry; import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint; import org.opensearch.node.remotestore.RemoteStorePinnedTimestampService; import org.opensearch.threadpool.ThreadPool; @@ -97,6 +98,8 @@ public final class RemoteSegmentStoreDirectory extends FilterDirectory implement private final ThreadPool threadPool; + private final ActiveMergesSegmentRegistry activeMergesSegmentRegistry = ActiveMergesSegmentRegistry.getInstance(); + /** * Keeps track of local segment filename to uploaded filename along with other attributes like checksum. * This map acts as a cache layer for uploaded segment filenames which helps avoid calling listAll() each time. @@ -354,6 +357,10 @@ public String getOriginalFilename() { return originalFilename; } + public String getUploadedFilename() { + return uploadedFilename; + } + public void setWrittenByMajor(int writtenByMajor) { if (writtenByMajor <= Version.LATEST.major && writtenByMajor >= Version.MIN_SUPPORTED_MAJOR) { this.writtenByMajor = writtenByMajor; @@ -369,6 +376,8 @@ public void setWrittenByMajor(int writtenByMajor) { ); } } + + public static UploadedSegmentMetadata EMPTY = new UploadedSegmentMetadata(null, null, null, -1); } /** @@ -500,12 +509,15 @@ public long fileLength(String name) throws IOException { if (segmentsUploadedToRemoteStore.containsKey(name)) { return segmentsUploadedToRemoteStore.get(name).getLength(); } + if (isRemoteStoreFileName(name)){ + return remoteDataDirectory.fileLength(name); + } String remoteFilename = getExistingRemoteFilename(name); if (remoteFilename != null) { return remoteDataDirectory.fileLength(remoteFilename); - } else { - throw new NoSuchFileException(name); } + throw new NoSuchFileException(name); + } /** @@ -530,7 +542,7 @@ public IndexOutput createOutput(String name, IOContext context) throws IOExcepti */ @Override public IndexInput openInput(String name, IOContext context) throws IOException { - String remoteFilename = getExistingRemoteFilename(name); + String remoteFilename = isRemoteStoreFileName(name) ? name : getExistingRemoteFilename(name); long fileLength = fileLength(name); if (remoteFilename != null) { return remoteDataDirectory.openInput(remoteFilename, fileLength, context); @@ -669,7 +681,53 @@ String getMetadataFileForCommit(long primaryTerm, long generation) throws IOExce private void postUpload(Directory from, String src, String remoteFilename, String checksum) throws IOException { UploadedSegmentMetadata segmentMetadata = new UploadedSegmentMetadata(src, remoteFilename, checksum, from.fileLength(src)); + if(activeMergesSegmentRegistry.contains(src)){ + activeMergesSegmentRegistry.updateMetadata(src, segmentMetadata); + return; + } segmentsUploadedToRemoteStore.put(src, segmentMetadata); + + } + + /** + * Synchronizes the {@link #segmentsUploadedToRemoteStore} map with the link{{@link ActiveMergesSegmentRegistry}} . + * For each segment in the provided collection: + * - If the segment exists in activeMergesSegmentRegistry, verifies the checksum matches + * - If checksums match + * - Copies the metadata from activeMergesSegmentRegistry to segmentsUploadedToRemoteStore + * - Unregisters the segment from activeMergesSegmentRegistry + * - If checksums don't match + * - Segment will be re-uploaded with a new name and old file will cleaned up by GC + * - Unregisters the segment from activeMergesSegmentRegistry + * - The segment file will automatically be deleted by the GC in time + * @param directory The directory containing the segment files + * @param segments Collection of segment names to synchronize + */ + public void syncSegmentsUploadedToRemoteStoreWithActiveMergesSegmentRegistry(Directory directory, Collection segments) { + segments.forEach(segment -> { + try { + if (activeMergesSegmentRegistry.contains(segment) == true) { + // TODO@kheraadi: Do we need to compare the checksum here? + String localChecksum = getChecksumOfLocalFile(directory, segment); + UploadedSegmentMetadata metadata = activeMergesSegmentRegistry.getMetadata(segment); + String storedChecksum = metadata.getChecksum(); + if (localChecksum.equals(storedChecksum)) { + segmentsUploadedToRemoteStore.put(segment, metadata); + } else { + // No-op, the segment file will be uploaded to the remote store again with + // a different name. GC will clean up the older file in time. + } + activeMergesSegmentRegistry.unregister(segment); + } + } catch (IOException e) { + logger.error("Exception while updating segmentsUploadedToRemoteStore for segment {}", segment, e); + } + }); + } + + private boolean isRemoteStoreFileName(String name) { + // TODO@kheraadi: Do we have a better way to check this? + return name.contains(SEGMENT_NAME_UUID_SEPARATOR); } /** @@ -988,6 +1046,7 @@ public void deleteStaleSegments(int lastNMetadataFilesToKeep) throws IOException staleSegmentRemoteFilenames.stream() .filter(file -> activeSegmentRemoteFilenames.contains(file) == false) .filter(file -> deletedSegmentFiles.contains(file) == false) + .filter(file -> activeMergesSegmentRegistry.canDelete(file) == true) .forEach(file -> { try { remoteDataDirectory.deleteFile(file); diff --git a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryFactory.java b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryFactory.java index 233665e65aed9..0b2d76d46772c 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryFactory.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryFactory.java @@ -85,7 +85,8 @@ public Directory newDirectory(String repositoryName, String indexUUID, ShardId s blobStoreRepository.blobStore().blobContainer(dataPath), blobStoreRepository::maybeRateLimitRemoteUploadTransfers, blobStoreRepository::maybeRateLimitLowPriorityRemoteUploadTransfers, - blobStoreRepository::maybeRateLimitRemoteDownloadTransfers + blobStoreRepository::maybeRateLimitRemoteDownloadTransfers, + blobStoreRepository::maybeRateLimitLowPriorityDownloadTransfers ); RemoteStorePathStrategy.ShardDataPathInput mdPathInput = RemoteStorePathStrategy.ShardDataPathInput.builder() diff --git a/server/src/main/java/org/opensearch/index/store/RemoteStoreFileDownloader.java b/server/src/main/java/org/opensearch/index/store/RemoteStoreFileDownloader.java index ad42b6d677b41..cfd4f92bc7ad6 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteStoreFileDownloader.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteStoreFileDownloader.java @@ -14,6 +14,7 @@ import org.opensearch.action.support.GroupedActionListener; import org.opensearch.action.support.PlainActionFuture; import org.opensearch.common.Nullable; +import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.annotation.PublicApi; import org.opensearch.common.logging.Loggers; import org.opensearch.common.util.CancellableThreads; @@ -41,6 +42,9 @@ public final class RemoteStoreFileDownloader { private final ThreadPool threadPool; private final RecoverySettings recoverySettings; + @ExperimentalApi + public record FileCopySpec (String localFilename, String remoteFilename, Long fileLength) {} + public RemoteStoreFileDownloader(ShardId shardId, ThreadPool threadPool, RecoverySettings recoverySettings) { this.logger = Loggers.getLogger(RemoteStoreFileDownloader.class, shardId); this.threadPool = threadPool; @@ -105,16 +109,46 @@ public void download( } } - private void downloadInternal( + // TODO@kheraadi: Fix this + public void download1( + Directory source, + Directory destination, + Directory secondDestination, + Collection toDownloadSegments, + Runnable onFileCompletion + ) throws InterruptedException, IOException { + final CancellableThreads cancellableThreads = new CancellableThreads(); + final PlainActionFuture listener = PlainActionFuture.newFuture(); + downloadInternal(cancellableThreads, source, destination, secondDestination, toDownloadSegments, onFileCompletion, listener); + try { + listener.get(); + } catch (ExecutionException e) { + if (e.getCause() instanceof RuntimeException) { + throw (RuntimeException) e.getCause(); + } else if (e.getCause() instanceof IOException) { + throw (IOException) e.getCause(); + } + throw new RuntimeException(e); + } catch (InterruptedException e) { + // If the blocking call on the PlainActionFuture itself is interrupted, then we must + // cancel the asynchronous work we were waiting on + cancellableThreads.cancel(e.getMessage()); + Thread.currentThread().interrupt(); + throw e; + } + } + + + private void downloadInternal( CancellableThreads cancellableThreads, Directory source, Directory destination, @Nullable Directory secondDestination, - Collection toDownloadSegments, + Collection toDownloadSegments, Runnable onFileCompletion, ActionListener listener ) { - final Queue queue = new ConcurrentLinkedQueue<>(toDownloadSegments); + final Queue queue = new ConcurrentLinkedQueue<>(toDownloadSegments); // Choose the minimum of: // - number of files to download // - max thread pool size @@ -130,29 +164,37 @@ private void downloadInternal( } } - private void copyOneFile( + private void copyOneFile( CancellableThreads cancellableThreads, Directory source, Directory destination, @Nullable Directory secondDestination, - Queue queue, + Queue queue, Runnable onFileCompletion, ActionListener listener ) { - final String file = queue.poll(); + final T file = queue.poll(); if (file == null) { // Queue is empty, so notify listener we are done listener.onResponse(null); } else { + String fileSrc, fileDest; + if (file instanceof FileCopySpec) { + FileCopySpec spec = ((FileCopySpec) file); + fileSrc = spec.remoteFilename(); + fileDest = spec.localFilename(); + } else { + fileSrc = fileDest = (String) file; + } threadPool.executor(ThreadPool.Names.REMOTE_RECOVERY).submit(() -> { logger.trace("Downloading file {}", file); try { cancellableThreads.executeIO(() -> { - destination.copyFrom(source, file, file, IOContext.DEFAULT); - logger.trace("Downloaded file {} of size {}", file, destination.fileLength(file)); + destination.copyFrom(source, fileSrc, fileDest, IOContext.DEFAULT); + logger.trace("Downloaded file {} of size {}", file, destination.fileLength(fileDest)); onFileCompletion.run(); if (secondDestination != null) { - secondDestination.copyFrom(destination, file, file, IOContext.DEFAULT); + secondDestination.copyFrom(destination, fileSrc, fileDest, IOContext.DEFAULT); } }); } catch (Exception e) { diff --git a/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java b/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java index 672a875f9ad00..01e8f19666520 100644 --- a/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java +++ b/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java @@ -8,24 +8,25 @@ package org.opensearch.indices.replication; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; +import org.opensearch.index.store.RemoteSegmentStoreDirectory.UploadedSegmentMetadata; import reactor.util.annotation.NonNull; +import java.util.Collections; import java.util.Map; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.locks.ReentrantLock; /** - * TODO@kheraadi: Simplify this - * Usages: - * 1. Used by GC to exclude merged segment files for deletion - * 2. Used to track remote store segment file names (remove this) - * - */ + * Registry to track active merge segments and their metadata. + * This class is implemented as a singleton and provides methods to register, update metadata, + * and unregister segment files that are being merged. It also maintains a mapping between + * local segment filenames and their corresponding uploaded segment metadata. + * */ public class ActiveMergesSegmentRegistry { - private final Map myMergedSegments = new ConcurrentHashMap<>(); - private static final String PENDING_REMOTE_FILE_NAME = "PENDING_REMOTE_FILE_NAME"; + private final Map segmentMetadataMap = new ConcurrentHashMap<>(); + private final Set filenameRegistry = ConcurrentHashMap.newKeySet(); + private final ReentrantLock lock = new ReentrantLock(); private static class HOLDER { private static final ActiveMergesSegmentRegistry INSTANCE = new ActiveMergesSegmentRegistry(); @@ -37,53 +38,78 @@ public static ActiveMergesSegmentRegistry getInstance() { return HOLDER.INSTANCE; } - public void updateRemoteSegmentFileName(@NonNull String localSegmentFileName, @NonNull String remoteSegmentFileName) { - if (contains(localSegmentFileName) && PENDING_REMOTE_FILE_NAME.equals(getExistingRemoteSegmentFileName(localSegmentFileName)) == false){ - // This should never happen - if (remoteSegmentFileName.equals(getExistingRemoteSegmentFileName(localSegmentFileName))) { - return; + /** + * Registers a segment file. Throws exception if already registered. + * @param localSegmentFilename Segment filename in the local store + */ + public void register(@NonNull String localSegmentFilename) { + lock.lock(); + try { + if (contains(localSegmentFilename)){ + throw new IllegalArgumentException(localSegmentFilename + " is already registered. Cannot reregister."); } - throw new IllegalArgumentException("Segment " + localSegmentFileName + " is already registered as " + getExistingRemoteSegmentFileName(localSegmentFileName) + ". Called with " + remoteSegmentFileName); + filenameRegistry.add(localSegmentFilename); + } finally { + lock.unlock(); } - myMergedSegments.put(localSegmentFileName, remoteSegmentFileName); } - public void register(@NonNull String localSegmentFileName) { - if (contains(localSegmentFileName)){ - if(PENDING_REMOTE_FILE_NAME.equals(getExistingRemoteSegmentFileName(localSegmentFileName))) { - return; + /** + * Adds {@link UploadedSegmentMetadata} for a registered segment file. Throws an error if the file is not registered. + * @param localSegmentFilename Segment filename in the local store + * @param metadata {@link UploadedSegmentMetadata} for the segment file + */ + public void updateMetadata(@NonNull String localSegmentFilename, @NonNull UploadedSegmentMetadata metadata) { + lock.lock(); + try { + if (contains(localSegmentFilename) == false) { + throw new IllegalArgumentException("Segment " + localSegmentFilename + " is not registered"); } - throw new IllegalArgumentException(localSegmentFileName + ": " + getExistingRemoteSegmentFileName(localSegmentFileName) + " already registered. Cannot reregister."); + segmentMetadataMap.put(localSegmentFilename, metadata); + filenameRegistry.add(metadata.getUploadedFilename()); + } finally { + lock.unlock(); } - - myMergedSegments.put(localSegmentFileName, PENDING_REMOTE_FILE_NAME); } - public void unregister(@NonNull String segmentFileName) { - myMergedSegments.remove(segmentFileName); + /** + * Unregisters a segment file from the registry. + * @param segmentFilename Segment filename in local store + */ + public void unregister(@NonNull String segmentFilename) { + lock.lock(); + try { + if (segmentMetadataMap.containsKey(segmentFilename)) { + String remoteFilename = segmentMetadataMap.get(segmentFilename).getUploadedFilename(); + filenameRegistry.remove(remoteFilename); + } + filenameRegistry.remove(segmentFilename); + segmentMetadataMap.remove(segmentFilename); + } finally { + lock.unlock(); + } } - public boolean contains(@NonNull String segmentFileName) { - return myMergedSegments.containsKey(segmentFileName); + public boolean contains(@NonNull String segmentFilename) { + return filenameRegistry.contains(segmentFilename); } - public String getExistingRemoteSegmentFileName(@NonNull String localSegmentFileName) { - if (contains(localSegmentFileName) == false) { - // This should never happen - throw new IllegalArgumentException("Segment " + localSegmentFileName + " is not registered"); + public String getExistingRemoteSegmentFilename(@NonNull String localSegmentFilename) { + if (segmentMetadataMap.containsKey(localSegmentFilename) == false) { + throw new IllegalArgumentException("Metadata for segment " + localSegmentFilename + " is not available."); } + return segmentMetadataMap.get(localSegmentFilename).getUploadedFilename(); + } - return myMergedSegments.get(localSegmentFileName); + public boolean canDelete(@NonNull String segmentFilename) { + return contains(segmentFilename) == false; } - public boolean canDelete(@NonNull String segmentFileName) { - String originalFileName = getOriginalFileName(segmentFileName); - return contains(originalFileName) && - segmentFileName.equals(getExistingRemoteSegmentFileName(originalFileName)); + public Map segmentMetadataMap() { + return Collections.unmodifiableMap(segmentMetadataMap); } - private String getOriginalFileName(@NonNull String remoteSegmentFileName) { - String originalFileName = remoteSegmentFileName.split("__")[0]; - return originalFileName; + public UploadedSegmentMetadata getMetadata(String localSegmentFilename){ + return segmentMetadataMap.get(localSegmentFilename); } } diff --git a/server/src/main/java/org/opensearch/indices/replication/RemoteStoreReplicationSource.java b/server/src/main/java/org/opensearch/indices/replication/RemoteStoreReplicationSource.java index bb79fab32cf9f..c8cf6d093816e 100644 --- a/server/src/main/java/org/opensearch/indices/replication/RemoteStoreReplicationSource.java +++ b/server/src/main/java/org/opensearch/indices/replication/RemoteStoreReplicationSource.java @@ -20,12 +20,15 @@ import org.opensearch.index.shard.IndexShard; import org.opensearch.index.shard.IndexShardState; import org.opensearch.index.store.RemoteSegmentStoreDirectory; +import org.opensearch.index.store.RemoteStoreFileDownloader; import org.opensearch.index.store.Store; import org.opensearch.index.store.StoreFileMetadata; import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadata; +import org.opensearch.indices.replication.checkpoint.RemoteStoreMergedSegmentCheckpoint; import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint; import java.io.IOException; +import java.rmi.Remote; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; @@ -154,6 +157,58 @@ public void getSegmentFiles( } } + /** + * {@inheritDoc} + * Refactor this + **/ + @Override + public void getMergedSegmentFiles( + long replicationId, + ReplicationCheckpoint checkpoint, + List filesToFetch, + IndexShard indexShard, + BiConsumer fileProgressTracker, + ActionListener listener + ) { + try { + assert checkpoint instanceof RemoteStoreMergedSegmentCheckpoint; + + if (filesToFetch.isEmpty()) { + listener.onResponse(new GetSegmentFilesResponse(Collections.emptyList())); + return; + } + + RemoteStoreMergedSegmentCheckpoint mergedSegmentCheckpoint = (RemoteStoreMergedSegmentCheckpoint) checkpoint; + + final Directory storeDirectory = indexShard.store().directory(); + + Map localToRemoteSegmentFileNameMap = mergedSegmentCheckpoint.getLocalToRemoteSegmentFilenameMap(); + + final List toDownloadSegmentNames = new ArrayList<>(); + for (StoreFileMetadata fileMetadata : filesToFetch) { + String file = fileMetadata.name(); + toDownloadSegmentNames.add( + new RemoteStoreFileDownloader.FileCopySpec( + file, + localToRemoteSegmentFileNameMap.get(file), + checkpoint.getMetadataMap().get(file).length()) + ); + } + + indexShard.getFileDownloader() + .download1( + indexShard.getRemoteDirectory(), + storeDirectory, + null, + toDownloadSegmentNames, + () -> listener.onResponse(new GetSegmentFilesResponse(filesToFetch)) + ); + + } catch (InterruptedException | IOException e) { + listener.onFailure(e); + } + } + @Override public void cancel() { this.cancellableThreads.cancel("Canceled by target"); diff --git a/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java b/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java index 1be0224a77b60..5f541486494ac 100644 --- a/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java +++ b/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java @@ -710,7 +710,7 @@ public void onReplicationFailure( } }); try { - latch.await(); + latch.await(); // @TODO@kheraadi: add a timeout here:: mergedSegmentTimeout - timeTakenToUpload } catch (InterruptedException e) { logger.warn( () -> new ParameterizedMessage("Interrupted while waiting for pre copy merged segment [{}]", receivedCheckpoint), diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/AbstractPublishCheckpointAction.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/AbstractPublishCheckpointAction.java index ddf7a1f61030c..0c76c34b697cf 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/AbstractPublishCheckpointAction.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/AbstractPublishCheckpointAction.java @@ -145,7 +145,7 @@ public String executor() { public void handleResponse(ReplicationResponse response) { try { timer.stop(); - logger.debug( + logger.info( () -> new ParameterizedMessage( "[shardId {}] Completed publishing checkpoint [{}], timing: {}", indexShard.shardId().getId(), @@ -164,7 +164,7 @@ public void handleResponse(ReplicationResponse response) { public void handleException(TransportException e) { try { timer.stop(); - logger.debug( + logger.info( "[shardId {}] Failed to publish checkpoint [{}], timing: {}", indexShard.shardId().getId(), checkpoint, @@ -213,12 +213,12 @@ public void handleException(TransportException e) { e ); } - } + } } } @Override - final protected void shardOperationOnReplica(ReplicaRequest shardRequest, IndexShard replica, ActionListener listener) { + protected void shardOperationOnReplica(ReplicaRequest shardRequest, IndexShard replica, ActionListener listener) { Objects.requireNonNull(shardRequest); Objects.requireNonNull(replica); ActionListener.completeWith(listener, () -> { diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentPublisher.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentPublisher.java index 4f58b10500d1e..32a5fd1c441bb 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentPublisher.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentPublisher.java @@ -29,7 +29,7 @@ public MergedSegmentPublisher(PublishAction publishAction) { this.publishAction = Objects.requireNonNull(publishAction); } - public void publish(IndexShard indexShard, MergedSegmentCheckpoint checkpoint) { + public void publish(IndexShard indexShard, ReplicationCheckpoint checkpoint) { publishAction.publish(indexShard, checkpoint); } diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java index 605dcf18e94d6..1e5fd75339cb3 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java @@ -8,6 +8,7 @@ package org.opensearch.indices.replication.checkpoint; +import org.opensearch.common.Nullable; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.settings.Settings; import org.opensearch.core.common.io.stream.StreamInput; @@ -17,6 +18,8 @@ import org.opensearch.index.store.StoreFileMetadata; import java.io.IOException; +import java.util.Collections; +import java.util.HashMap; import java.util.Map; import java.util.Objects; @@ -28,7 +31,7 @@ @ExperimentalApi public class RemoteStoreMergedSegmentCheckpoint extends ReplicationCheckpoint { private final String segmentName; - private final String remoteStorePath; + private final Map localToRemoteSegmentFilenameMap; public RemoteStoreMergedSegmentCheckpoint( ShardId shardId, @@ -37,32 +40,24 @@ public RemoteStoreMergedSegmentCheckpoint( String codec, Map metadataMap, String segmentName, - String remoteStorePath + @Nullable Map localToRemoteSegmentFilenameMap ) { super(shardId, primaryTerm, SequenceNumbers.NO_OPS_PERFORMED, SequenceNumbers.NO_OPS_PERFORMED, length, codec, metadataMap); this.segmentName = segmentName; - this.remoteStorePath = remoteStorePath; + this.localToRemoteSegmentFilenameMap = localToRemoteSegmentFilenameMap == null ? new HashMap<>() : localToRemoteSegmentFilenameMap; } public RemoteStoreMergedSegmentCheckpoint(StreamInput in) throws IOException { super(in); - segmentName = in.readString(); - remoteStorePath = in.readString(); - } - - public String getSegmentName() { - return segmentName; - } - - public String getRemoteStorePath() { - return remoteStorePath; + this.segmentName = in.readString(); + this.localToRemoteSegmentFilenameMap = in.readMap(StreamInput::readString, StreamInput::readString); } @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); out.writeString(segmentName); - out.writeString(remoteStorePath); + out.writeMap(getLocalToRemoteSegmentFilenameMap(), StreamOutput::writeString, StreamOutput::writeString); } @Override @@ -72,14 +67,13 @@ public boolean equals(Object o) { RemoteStoreMergedSegmentCheckpoint that = (RemoteStoreMergedSegmentCheckpoint) o; return getPrimaryTerm() == that.getPrimaryTerm() && segmentName.equals(that.segmentName) - && remoteStorePath.equals(that.remoteStorePath) && Objects.equals(getShardId(), that.getShardId()) && getCodec().equals(that.getCodec()); } @Override public int hashCode() { - return Objects.hash(getShardId(), getPrimaryTerm(), segmentName, remoteStorePath); + return Objects.hash(getShardId(), getPrimaryTerm(), segmentName); } @Override @@ -90,13 +84,21 @@ public String toString() { + ", primaryTerm=" + getPrimaryTerm() + ", segmentName=" - + segmentName - + ", remoteStorePath=" - + remoteStorePath + + getSegmentName() + + ", localToRemoteSegmentFilenameSize=" + + getLocalToRemoteSegmentFilenameMap().size() + '}'; } - public Map getLocalToRemoteSegmentFileNameMap() { - return null; + public Map getLocalToRemoteSegmentFilenameMap() { + return this.localToRemoteSegmentFilenameMap; + } + + public String getSegmentName() { + return segmentName; + } + + public void updateLocalToRemoteSegmentFilenameMap(String localSegmentFilename, String remoteSegmentFilename) { + localToRemoteSegmentFilenameMap.put(localSegmentFilename, remoteSegmentFilename); } } diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java index 9c2d64fd90008..807bba0f7a768 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java @@ -10,16 +10,19 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.apache.logging.log4j.message.ParameterizedMessage; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.replication.ReplicationResponse; import org.opensearch.cluster.action.shard.ShardStateAction; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.UploadListener; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.index.shard.IndexShard; +import org.opensearch.index.shard.RemoteStoreUploaderService; import org.opensearch.indices.IndicesService; import org.opensearch.indices.replication.ActiveMergesSegmentRegistry; import org.opensearch.indices.replication.SegmentReplicationTargetService; @@ -28,6 +31,7 @@ import java.util.Collection; import java.util.Map; +import java.util.Objects; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; @@ -72,39 +76,35 @@ public RemoteStorePublishMergedSegmentAction( @Override protected void doReplicaOperation(RemoteStorePublishMergedSegmentRequest shardRequest, IndexShard replica) { - + if (shardRequest.getMergedSegment().getShardId().equals(replica.shardId())) { + replicationService.onNewMergedSegmentCheckpoint(shardRequest.getMergedSegment(), replica); + } } @Override protected void shardOperationOnPrimary(RemoteStorePublishMergedSegmentRequest shardRequest, IndexShard primary, ActionListener> listener) { - + ActionListener.completeWith(listener, () -> new PrimaryResult<>(shardRequest, new ReplicationResponse())); } + @Override public void publish(IndexShard indexShard, ReplicationCheckpoint checkpoint) { assert checkpoint instanceof RemoteStoreMergedSegmentCheckpoint; RemoteStoreMergedSegmentCheckpoint mergedSegmentCheckpoint = (RemoteStoreMergedSegmentCheckpoint) checkpoint; - publishMergedSegmentsToRemoteStore(indexShard, mergedSegmentCheckpoint); - logger.info("RemoteFileNames at {}#publish {}", getClass().getName(), mergedSegmentCheckpoint).getLocalToRemoteSegmentFileNameMap()); doPublish(indexShard, checkpoint, - new RemoteStorePublishMergedSegmentRequest((RemoteStoreMergedSegmentCheckpoint) checkpoint), + new RemoteStorePublishMergedSegmentRequest(mergedSegmentCheckpoint), "segrep_publish_merged_segment", true, indexShard.getRecoverySettings().getMergedSegmentReplicationTimeout() ); - mergedSegmentCheckpoint - .getLocalToRemoteSegmentFileNameMap() - .keySet() - .forEach(activeMergesSegmentRegistry::unregister); } private void publishMergedSegmentsToRemoteStore(IndexShard indexShard, RemoteStoreMergedSegmentCheckpoint checkpoint) { RemoteStoreUploaderService remoteStoreUploaderService = getRemoteStoreUploaderService(indexShard); Collection segmentsToUpload = checkpoint.getMetadataMap().keySet(); - logger.info("Publishing segments {} to remote store", segmentsToUpload); - registerSegmentsToActiveMerges(segmentsToUpload); + Map segmentsSizeMap = checkpoint .getMetadataMap() .entrySet() @@ -113,13 +113,16 @@ private void publishMergedSegmentsToRemoteStore(IndexShard indexShard, RemoteSto final CountDownLatch latch = new CountDownLatch(segmentsToUpload.size()); + // TODO@kheraadi: Upload in low priority remoteStoreUploaderService.uploadSegments( segmentsToUpload, segmentsSizeMap, new ActionListener() { @Override public void onResponse(Void unused) { - logger.info("Successfully uploaded segments {} to remote store", checkpoint.getLocalToRemoteSegmentFileNameMap()); + if(logger.isTraceEnabled() == true) { + logger.trace("Successfully uploaded segments {} to remote store", segmentsToUpload); + } } @Override @@ -132,15 +135,14 @@ public void onFailure(Exception e) { (x) -> new UploadListener() { @Override public void beforeUpload(String file) { - + activeMergesSegmentRegistry.register(file); } @Override public void onSuccess(String file) { - logger.info("Uploaded {}", file); - checkpoint.addRemoteSegmentFileName( + checkpoint.updateLocalToRemoteSegmentFilenameMap( file, - activeMergesSegmentRegistry.getExistingRemoteSegmentFileName(file) + activeMergesSegmentRegistry.getExistingRemoteSegmentFilename(file) ); latch.countDown(); } @@ -174,11 +176,4 @@ private RemoteStoreUploaderService getRemoteStoreUploaderService(IndexShard inde indexShard.getRemoteDirectory() ); } - - private void registerSegmentsToActiveMerges(Collection segmentsToUpload) { - logger.info("Registering segments to active merges: {}", segmentsToUpload); - segmentsToUpload.forEach(activeMergesSegmentRegistry::register); - logger.info("ActiveMergesSegmentsRegistry: " + activeMergesSegmentsRegistry.mergedSegments().toString()); - } - } diff --git a/server/src/main/java/org/opensearch/repositories/FilterRepository.java b/server/src/main/java/org/opensearch/repositories/FilterRepository.java index 188d61e00f206..b4bc88b6c9914 100644 --- a/server/src/main/java/org/opensearch/repositories/FilterRepository.java +++ b/server/src/main/java/org/opensearch/repositories/FilterRepository.java @@ -176,6 +176,11 @@ public long getRemoteDownloadThrottleTimeInNanos() { return in.getRemoteDownloadThrottleTimeInNanos(); } + @Override + public long getLowPriorityRemoteDownloadThrottleTimeInNanos() { + return in.getRemoteDownloadThrottleTimeInNanos(); + } + @Override public String startVerification() { return in.startVerification(); diff --git a/server/src/main/java/org/opensearch/repositories/Repository.java b/server/src/main/java/org/opensearch/repositories/Repository.java index 259c4a6e09ce7..7a6ada3c755ed 100644 --- a/server/src/main/java/org/opensearch/repositories/Repository.java +++ b/server/src/main/java/org/opensearch/repositories/Repository.java @@ -302,6 +302,11 @@ default long getLowPriorityRemoteUploadThrottleTimeInNanos() { */ long getRemoteDownloadThrottleTimeInNanos(); + /** + * Returns low priority remote download throttle time in nanoseconds + */ + long getLowPriorityRemoteDownloadThrottleTimeInNanos(); + /** * Returns stats on the repository usage */ diff --git a/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java b/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java index 940942b816536..95c889fc50f17 100644 --- a/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java +++ b/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java @@ -444,12 +444,16 @@ protected static long calculateMaxWithinIntLimit(long defaultThresholdOfHeap, lo private volatile RateLimiter remoteDownloadRateLimiter; + private volatile RateLimiter remoteDownloadLowPriorityRateLimiter; + private final CounterMetric snapshotRateLimitingTimeInNanos = new CounterMetric(); private final CounterMetric restoreRateLimitingTimeInNanos = new CounterMetric(); private final CounterMetric remoteDownloadRateLimitingTimeInNanos = new CounterMetric(); + private final CounterMetric remoteDownloadLowPriorityRateLimitingTimeInNanos = new CounterMetric(); + private final CounterMetric remoteUploadRateLimitingTimeInNanos = new CounterMetric(); private final CounterMetric remoteUploadLowPriorityRateLimitingTimeInNanos = new CounterMetric(); @@ -601,6 +605,7 @@ private void readRepositoryMetadata(RepositoryMetadata repositoryMetadata) { ByteSizeValue.ZERO ); remoteDownloadRateLimiter = getRateLimiter(metadata.settings(), "max_remote_download_bytes_per_sec", ByteSizeValue.ZERO); + remoteDownloadLowPriorityRateLimiter = getRateLimiter(metadata.settings(), "max_remote_download_bytes_per_sec", ByteSizeValue.ZERO); readOnly = READONLY_SETTING.get(metadata.settings()); cacheRepositoryData = CACHE_REPOSITORY_DATA.get(metadata.settings()); bufferSize = Math.toIntExact(BUFFER_SIZE_SETTING.get(metadata.settings()).getBytes()); @@ -2930,6 +2935,11 @@ public long getRemoteDownloadThrottleTimeInNanos() { return remoteDownloadRateLimitingTimeInNanos.count(); } + @Override + public long getLowPriorityRemoteDownloadThrottleTimeInNanos() { + return remoteDownloadLowPriorityRateLimitingTimeInNanos.count(); + } + protected void assertSnapshotOrGenericThread() { assert Thread.currentThread().getName().contains('[' + ThreadPool.Names.SNAPSHOT_DELETION + ']') || Thread.currentThread().getName().contains('[' + ThreadPool.Names.SNAPSHOT + ']') @@ -4322,6 +4332,20 @@ public InputStream maybeRateLimitRemoteDownloadTransfers(InputStream inputStream ); } + public InputStream maybeRateLimitLowPriorityDownloadTransfers(InputStream inputStream) { + return maybeRateLimit( + maybeRateLimit( + inputStream, + () -> remoteDownloadLowPriorityRateLimiter, + remoteDownloadLowPriorityRateLimitingTimeInNanos, + BlobStoreTransferContext.REMOTE_DOWNLOAD + ), + recoverySettings::recoveryRateLimiter, + remoteDownloadRateLimitingTimeInNanos, + BlobStoreTransferContext.REMOTE_DOWNLOAD + ); + } + public InputStream maybeRateLimitSnapshots(InputStream stream) { return maybeRateLimit(stream, () -> snapshotRateLimiter, snapshotRateLimitingTimeInNanos, BlobStoreTransferContext.SNAPSHOT); } From 6d62b6894c4c3c435b3dc00eff19fa058e1cd28b Mon Sep 17 00:00:00 2001 From: kh3ra Date: Wed, 2 Jul 2025 18:50:32 +0530 Subject: [PATCH 5/6] minor changes and addressed some TODOs --- .../store/RemoteSegmentStoreDirectory.java | 1 - .../store/RemoteStoreFileDownloader.java | 45 +++---------------- .../RemoteStoreReplicationSource.java | 6 +-- .../AbstractPublishCheckpointAction.java | 4 +- .../checkpoint/MergedSegmentCheckpoint.java | 2 +- .../PublishMergedSegmentActionProvider.java | 8 +++- .../RemoteStoreMergedSegmentCheckpoint.java | 17 +------ ...RemoteStorePublishMergedSegmentAction.java | 8 ++-- .../repositories/FilterRepository.java | 2 +- 9 files changed, 27 insertions(+), 66 deletions(-) diff --git a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java index 5cc27f55d9892..51c18d40722cb 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java @@ -707,7 +707,6 @@ public void syncSegmentsUploadedToRemoteStoreWithActiveMergesSegmentRegistry(Dir segments.forEach(segment -> { try { if (activeMergesSegmentRegistry.contains(segment) == true) { - // TODO@kheraadi: Do we need to compare the checksum here? String localChecksum = getChecksumOfLocalFile(directory, segment); UploadedSegmentMetadata metadata = activeMergesSegmentRegistry.getMetadata(segment); String storedChecksum = metadata.getChecksum(); diff --git a/server/src/main/java/org/opensearch/index/store/RemoteStoreFileDownloader.java b/server/src/main/java/org/opensearch/index/store/RemoteStoreFileDownloader.java index cfd4f92bc7ad6..899340026b7a6 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteStoreFileDownloader.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteStoreFileDownloader.java @@ -43,7 +43,7 @@ public final class RemoteStoreFileDownloader { private final RecoverySettings recoverySettings; @ExperimentalApi - public record FileCopySpec (String localFilename, String remoteFilename, Long fileLength) {} + public record FileCopySpec (String localFilename, String remoteFilename) {} public RemoteStoreFileDownloader(ShardId shardId, ThreadPool threadPool, RecoverySettings recoverySettings) { this.logger = Loggers.getLogger(RemoteStoreFileDownloader.class, shardId); @@ -76,45 +76,16 @@ public void downloadAsync( * @param destination The local directory to copy segment files to * @param secondDestination The second remote directory that segment files are * copied to after being copied to the local directory - * @param toDownloadSegments The list of segment files to download + * @param toDownloadSegments The list of segment files to download as String or {@link FileCopySpec} * @param onFileCompletion A generic runnable that is invoked after each file download. * Must be thread safe as this may be invoked concurrently from * different threads. */ - public void download( - Directory source, - Directory destination, - Directory secondDestination, - Collection toDownloadSegments, - Runnable onFileCompletion - ) throws InterruptedException, IOException { - final CancellableThreads cancellableThreads = new CancellableThreads(); - final PlainActionFuture listener = PlainActionFuture.newFuture(); - downloadInternal(cancellableThreads, source, destination, secondDestination, toDownloadSegments, onFileCompletion, listener); - try { - listener.get(); - } catch (ExecutionException e) { - if (e.getCause() instanceof RuntimeException) { - throw (RuntimeException) e.getCause(); - } else if (e.getCause() instanceof IOException) { - throw (IOException) e.getCause(); - } - throw new RuntimeException(e); - } catch (InterruptedException e) { - // If the blocking call on the PlainActionFuture itself is interrupted, then we must - // cancel the asynchronous work we were waiting on - cancellableThreads.cancel(e.getMessage()); - Thread.currentThread().interrupt(); - throw e; - } - } - - // TODO@kheraadi: Fix this - public void download1( + public void download( Directory source, Directory destination, Directory secondDestination, - Collection toDownloadSegments, + Collection toDownloadSegments, Runnable onFileCompletion ) throws InterruptedException, IOException { final CancellableThreads cancellableThreads = new CancellableThreads(); @@ -138,7 +109,6 @@ public void download1( } } - private void downloadInternal( CancellableThreads cancellableThreads, Directory source, @@ -179,10 +149,9 @@ private void copyOneFile( listener.onResponse(null); } else { String fileSrc, fileDest; - if (file instanceof FileCopySpec) { - FileCopySpec spec = ((FileCopySpec) file); - fileSrc = spec.remoteFilename(); - fileDest = spec.localFilename(); + if (file instanceof FileCopySpec(String localFilename, String remoteFilename)) { + fileSrc = remoteFilename; + fileDest = localFilename; } else { fileSrc = fileDest = (String) file; } diff --git a/server/src/main/java/org/opensearch/indices/replication/RemoteStoreReplicationSource.java b/server/src/main/java/org/opensearch/indices/replication/RemoteStoreReplicationSource.java index c8cf6d093816e..88389338a0bf4 100644 --- a/server/src/main/java/org/opensearch/indices/replication/RemoteStoreReplicationSource.java +++ b/server/src/main/java/org/opensearch/indices/replication/RemoteStoreReplicationSource.java @@ -190,13 +190,13 @@ public void getMergedSegmentFiles( toDownloadSegmentNames.add( new RemoteStoreFileDownloader.FileCopySpec( file, - localToRemoteSegmentFileNameMap.get(file), - checkpoint.getMetadataMap().get(file).length()) + localToRemoteSegmentFileNameMap.get(file) + ) ); } indexShard.getFileDownloader() - .download1( + .download( indexShard.getRemoteDirectory(), storeDirectory, null, diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/AbstractPublishCheckpointAction.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/AbstractPublishCheckpointAction.java index 0c76c34b697cf..9d9fb4d3b838d 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/AbstractPublishCheckpointAction.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/AbstractPublishCheckpointAction.java @@ -145,7 +145,7 @@ public String executor() { public void handleResponse(ReplicationResponse response) { try { timer.stop(); - logger.info( + logger.trace( () -> new ParameterizedMessage( "[shardId {}] Completed publishing checkpoint [{}], timing: {}", indexShard.shardId().getId(), @@ -164,7 +164,7 @@ public void handleResponse(ReplicationResponse response) { public void handleException(TransportException e) { try { timer.stop(); - logger.info( + logger.trace( "[shardId {}] Failed to publish checkpoint [{}], timing: {}", indexShard.shardId().getId(), checkpoint, diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentCheckpoint.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentCheckpoint.java index 6b8c82cafea8a..8bc1b596100ea 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentCheckpoint.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentCheckpoint.java @@ -51,7 +51,7 @@ */ @ExperimentalApi public class MergedSegmentCheckpoint extends ReplicationCheckpoint { - private final String segmentName; + final String segmentName; public MergedSegmentCheckpoint( ShardId shardId, diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java index 49f033e0264d9..ede2a276b6c90 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java @@ -10,6 +10,7 @@ import org.opensearch.common.inject.Inject; import org.opensearch.common.inject.Provider; import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.indices.IndicesService; import org.opensearch.indices.replication.SegmentReplicationTargetService; import org.opensearch.threadpool.ThreadPool; @@ -49,8 +50,11 @@ public PublishMergedSegmentActionProvider( @Override public MergedSegmentPublisher.PublishAction get() { - // TODO@kheraadi: check if remote store is enabled here - if (false) { + if (FeatureFlags.isEnabled(FeatureFlags.MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING) == false) { + return null; + } + // TODO@kheraadi: FIX THIS + if (false) {// || clusterService.localNode().isRemoteStoreNode() == false) { return new PublishMergedSegmentAction( settings, transportService, clusterService, indicesService, threadPool, shardStateAction, actionFilters, targetService diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java index 1e5fd75339cb3..dd430022e71ce 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java @@ -29,8 +29,7 @@ * @opensearch.internal */ @ExperimentalApi -public class RemoteStoreMergedSegmentCheckpoint extends ReplicationCheckpoint { - private final String segmentName; +public class RemoteStoreMergedSegmentCheckpoint extends MergedSegmentCheckpoint { private final Map localToRemoteSegmentFilenameMap; public RemoteStoreMergedSegmentCheckpoint( @@ -42,21 +41,18 @@ public RemoteStoreMergedSegmentCheckpoint( String segmentName, @Nullable Map localToRemoteSegmentFilenameMap ) { - super(shardId, primaryTerm, SequenceNumbers.NO_OPS_PERFORMED, SequenceNumbers.NO_OPS_PERFORMED, length, codec, metadataMap); - this.segmentName = segmentName; + super(shardId, primaryTerm, length, codec, metadataMap, segmentName); this.localToRemoteSegmentFilenameMap = localToRemoteSegmentFilenameMap == null ? new HashMap<>() : localToRemoteSegmentFilenameMap; } public RemoteStoreMergedSegmentCheckpoint(StreamInput in) throws IOException { super(in); - this.segmentName = in.readString(); this.localToRemoteSegmentFilenameMap = in.readMap(StreamInput::readString, StreamInput::readString); } @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); - out.writeString(segmentName); out.writeMap(getLocalToRemoteSegmentFilenameMap(), StreamOutput::writeString, StreamOutput::writeString); } @@ -71,11 +67,6 @@ public boolean equals(Object o) { && getCodec().equals(that.getCodec()); } - @Override - public int hashCode() { - return Objects.hash(getShardId(), getPrimaryTerm(), segmentName); - } - @Override public String toString() { return "RemoteStoreMergedSegmentCheckpoint{" @@ -94,10 +85,6 @@ public Map getLocalToRemoteSegmentFilenameMap() { return this.localToRemoteSegmentFilenameMap; } - public String getSegmentName() { - return segmentName; - } - public void updateLocalToRemoteSegmentFilenameMap(String localSegmentFilename, String remoteSegmentFilename) { localToRemoteSegmentFilenameMap.put(localSegmentFilename, remoteSegmentFilename); } diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java index 807bba0f7a768..6550cf645fa21 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java @@ -38,7 +38,7 @@ public class RemoteStorePublishMergedSegmentAction extends AbstractPublishCheckpointAction implements MergedSegmentPublisher.PublishAction { - public static final String ACTION_NAME = "indices:admin/publish_merged_segment"; + public static final String ACTION_NAME = "indices:admin/remote_publish_merged_segment"; private final static Logger logger = LogManager.getLogger(RemoteStorePublishMergedSegmentAction.class); @@ -89,8 +89,10 @@ protected void shardOperationOnPrimary(RemoteStorePublishMergedSegmentRequest sh @Override public void publish(IndexShard indexShard, ReplicationCheckpoint checkpoint) { - assert checkpoint instanceof RemoteStoreMergedSegmentCheckpoint; - RemoteStoreMergedSegmentCheckpoint mergedSegmentCheckpoint = (RemoteStoreMergedSegmentCheckpoint) checkpoint; + if(! (checkpoint instanceof RemoteStoreMergedSegmentCheckpoint mergedSegmentCheckpoint)) { + throw new AssertionError("Expected checkpoint to be an instance of " + RemoteStoreMergedSegmentCheckpoint.class); + } + publishMergedSegmentsToRemoteStore(indexShard, mergedSegmentCheckpoint); doPublish(indexShard, checkpoint, diff --git a/server/src/main/java/org/opensearch/repositories/FilterRepository.java b/server/src/main/java/org/opensearch/repositories/FilterRepository.java index b4bc88b6c9914..60b1ab9cd871a 100644 --- a/server/src/main/java/org/opensearch/repositories/FilterRepository.java +++ b/server/src/main/java/org/opensearch/repositories/FilterRepository.java @@ -178,7 +178,7 @@ public long getRemoteDownloadThrottleTimeInNanos() { @Override public long getLowPriorityRemoteDownloadThrottleTimeInNanos() { - return in.getRemoteDownloadThrottleTimeInNanos(); + return in.getLowPriorityRemoteDownloadThrottleTimeInNanos(); } @Override From 839a7f5ff9d45748c57a2a461864ac752438cb0b Mon Sep 17 00:00:00 2001 From: kh3ra Date: Fri, 4 Jul 2025 12:53:59 +0530 Subject: [PATCH 6/6] Adding unit tests --- .../store/RemoteSegmentStoreDirectory.java | 2 +- .../ActiveMergesSegmentRegistry.java | 2 +- .../PublishMergedSegmentActionProvider.java | 2 +- ...RemoteStorePublishMergedSegmentAction.java | 12 +- .../ActiveMergesSegmentRegistryTests.java | 192 +++++++++++ .../MergedSegmentReplicationTargetTests.java | 54 ++- ...eStorePublishMergedSegmentActionTests.java | 317 ++++++++++++++++++ ...StorePublishMergedSegmentRequestTests.java | 67 ++++ .../RepositoriesServiceTests.java | 5 + .../index/shard/RestoreOnlyRepository.java | 6 + 10 files changed, 640 insertions(+), 19 deletions(-) create mode 100644 server/src/test/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistryTests.java create mode 100644 server/src/test/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentActionTests.java create mode 100644 server/src/test/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentRequestTests.java diff --git a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java index 51c18d40722cb..ded50f355e901 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java @@ -725,7 +725,7 @@ public void syncSegmentsUploadedToRemoteStoreWithActiveMergesSegmentRegistry(Dir } private boolean isRemoteStoreFileName(String name) { - // TODO@kheraadi: Do we have a better way to check this? + // TODO: Do we have a better way to check this? return name.contains(SEGMENT_NAME_UUID_SEPARATOR); } diff --git a/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java b/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java index 01e8f19666520..2e3d2e0868d78 100644 --- a/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java +++ b/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java @@ -25,7 +25,7 @@ * */ public class ActiveMergesSegmentRegistry { private final Map segmentMetadataMap = new ConcurrentHashMap<>(); - private final Set filenameRegistry = ConcurrentHashMap.newKeySet(); + final Set filenameRegistry = ConcurrentHashMap.newKeySet(); // package-private for tests private final ReentrantLock lock = new ReentrantLock(); private static class HOLDER { diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java index ede2a276b6c90..145d96e1aee6c 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java @@ -53,7 +53,7 @@ public MergedSegmentPublisher.PublishAction get() { if (FeatureFlags.isEnabled(FeatureFlags.MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING) == false) { return null; } - // TODO@kheraadi: FIX THIS + // TODO: FIX THIS if (false) {// || clusterService.localNode().isRemoteStoreNode() == false) { return new PublishMergedSegmentAction( settings, transportService, clusterService, indicesService, diff --git a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java index 6550cf645fa21..e9fc61c6470fb 100644 --- a/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java @@ -115,7 +115,7 @@ private void publishMergedSegmentsToRemoteStore(IndexShard indexShard, RemoteSto final CountDownLatch latch = new CountDownLatch(segmentsToUpload.size()); - // TODO@kheraadi: Upload in low priority + // TODO: Upload in low priority remoteStoreUploaderService.uploadSegments( segmentsToUpload, segmentsSizeMap, @@ -153,23 +153,21 @@ public void onSuccess(String file) { public void onFailure(String file) { segmentsToUpload.forEach(activeMergesSegmentRegistry::unregister); /** - * TODO@kheraadi: - * 1. reset ActiveMergesRegistry - * 2. abort merge + * TODO: abort merge */ } } ); try { - if(latch.await(60, TimeUnit.MINUTES) == false) {throw new RuntimeException("Merged segment upload timed out.");}; // TODO@kheraadi: Finalize timeout + if(latch.await(60, TimeUnit.MINUTES) == false) {throw new RuntimeException("Merged segment upload timed out.");}; // TODO: Finalize timeout } catch (InterruptedException e) { throw new RuntimeException(e); - // TODO@kheraadi: abort merge properly here + // TODO: abort merge properly here } } /** - * TODO@kheraadi: REBASE ONCE UPLOAD CHANGES ARE COMPLETE + * TODO: REBASE ONCE UPLOAD CHANGES ARE COMPLETE */ private RemoteStoreUploaderService getRemoteStoreUploaderService(IndexShard indexShard) { return new RemoteStoreUploaderService( diff --git a/server/src/test/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistryTests.java b/server/src/test/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistryTests.java new file mode 100644 index 0000000000000..83aaeeeb46d0a --- /dev/null +++ b/server/src/test/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistryTests.java @@ -0,0 +1,192 @@ +package org.opensearch.indices.replication; + +import org.junit.After; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.OrderWith; +import org.opensearch.index.store.RemoteSegmentStoreDirectory.UploadedSegmentMetadata; + +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.*; +import static org.mockito.Mockito.*; + +public class ActiveMergesSegmentRegistryTests { + + private ActiveMergesSegmentRegistry registry; + private UploadedSegmentMetadata mockMetadata; + + @Before + public void setUp() { + registry = ActiveMergesSegmentRegistry.getInstance(); + // Clear registry state before each test + clearRegistry(); + + mockMetadata = mock(UploadedSegmentMetadata.class); + when(mockMetadata.getUploadedFilename()).thenReturn("remote_segment_1.si"); + } + + private void clearRegistry() { + // Clear all registered segments + Map metadataMap = registry.segmentMetadataMap(); + metadataMap.keySet().forEach(registry::unregister); + registry.filenameRegistry.clear(); + } + + @Test + public void testSingletonInstance() { + ActiveMergesSegmentRegistry instance1 = ActiveMergesSegmentRegistry.getInstance(); + ActiveMergesSegmentRegistry instance2 = ActiveMergesSegmentRegistry.getInstance(); + assertSame(instance1, instance2); + } + + @Test + public void testRegisterSegment() { + String filename = "segment_1.si"; + registry.register(filename); + assertTrue(registry.contains(filename)); + } + + @Test(expected = IllegalArgumentException.class) + public void testRegisterDuplicateSegment() { + String filename = "segment_1.si"; + registry.register(filename); + registry.register(filename); // Should throw exception + } + + @Test + public void testUpdateMetadata() { + String filename = "segment_1.si"; + registry.register(filename); + registry.updateMetadata(filename, mockMetadata); + + assertEquals(mockMetadata, registry.getMetadata(filename)); + assertTrue(registry.contains("remote_segment_1.si")); + } + + @Test(expected = IllegalArgumentException.class) + public void testUpdateMetadataUnregisteredSegment() { + registry.updateMetadata("unregistered_segment.si", mockMetadata); + } + + @Test + public void testUnregisterSegment() { + String filename = "segment_1.si"; + registry.register(filename); + registry.updateMetadata(filename, mockMetadata); + + registry.unregister(filename); + + assertFalse(registry.contains(filename)); + assertFalse(registry.contains("remote_segment_1.si")); + assertNull(registry.getMetadata(filename)); + } + + @Test + public void testUnregisterNonExistentSegment() { + // Should not throw exception + registry.unregister("non_existent.si"); + } + + @Test + public void testGetExistingRemoteSegmentFilename() { + String filename = "segment_1.si"; + registry.register(filename); + registry.updateMetadata(filename, mockMetadata); + + assertEquals("remote_segment_1.si", registry.getExistingRemoteSegmentFilename(filename)); + } + + @Test(expected = IllegalArgumentException.class) + public void testGetExistingRemoteSegmentFilenameNoMetadata() { + String filename = "segment_1.si"; + registry.register(filename); + registry.getExistingRemoteSegmentFilename(filename); // Metadata not available + } + + @Test + public void testCanDelete() { + String filename = "segment_1.si"; + assertTrue(registry.canDelete(filename)); // Not registered + + registry.register(filename); + assertFalse(registry.canDelete(filename)); // Registered + + registry.unregister(filename); + assertTrue(registry.canDelete(filename)); // Unregistered + } + + @Test + public void testConcurrentAccess() throws InterruptedException { + int threadCount = 10; + int operationsPerThread = 100; + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + CountDownLatch latch = new CountDownLatch(threadCount); + + for (int i = 0; i < threadCount; i++) { + final int threadId = i; + executor.submit(() -> { + try { + for (int j = 0; j < operationsPerThread; j++) { + String filename = "segment_" + threadId + "_" + j + ".si"; + String remoteFilename = "remote_" + filename; + UploadedSegmentMetadata metadata = mock(UploadedSegmentMetadata.class); + when(metadata.getUploadedFilename()).thenReturn(remoteFilename); + registry.register(filename); + assertTrue(registry.contains(filename)); + registry.updateMetadata(filename, metadata); + assertEquals(registry.getExistingRemoteSegmentFilename(filename), remoteFilename); + registry.unregister(filename); + assertFalse(registry.contains(filename)); + } + } finally { + latch.countDown(); + } + }); + } + + assertTrue(latch.await(30, TimeUnit.SECONDS)); + executor.shutdown(); + } + + @Test + public void testMultipleSegmentsLifecycle() { + String[] filenames = {"seg1.si", "seg2.si", "seg3.si"}; + UploadedSegmentMetadata[] metadatas = new UploadedSegmentMetadata[3]; + + // Setup mocks + for (int i = 0; i < 3; i++) { + metadatas[i] = mock(UploadedSegmentMetadata.class); + when(metadatas[i].getUploadedFilename()).thenReturn("remote_" + filenames[i]); + } + + // Register all + for (String filename : filenames) { + registry.register(filename); + assertTrue(registry.contains(filename)); + } + + // Update metadata + for (int i = 0; i < 3; i++) { + registry.updateMetadata(filenames[i], metadatas[i]); + assertEquals(metadatas[i], registry.getMetadata(filenames[i])); + } + + // Verify all are tracked + assertEquals(3, registry.segmentMetadataMap().size()); + + // Unregister one + registry.unregister(filenames[1]); + assertFalse(registry.contains(filenames[1])); + assertEquals(2, registry.segmentMetadataMap().size()); + + // Others still exist + assertTrue(registry.contains(filenames[0])); + assertTrue(registry.contains(filenames[2])); + } +} diff --git a/server/src/test/java/org/opensearch/indices/replication/MergedSegmentReplicationTargetTests.java b/server/src/test/java/org/opensearch/indices/replication/MergedSegmentReplicationTargetTests.java index 1d9e47a7fc2f4..d11de7f9c06f9 100644 --- a/server/src/test/java/org/opensearch/indices/replication/MergedSegmentReplicationTargetTests.java +++ b/server/src/test/java/org/opensearch/indices/replication/MergedSegmentReplicationTargetTests.java @@ -13,6 +13,7 @@ import org.apache.lucene.store.ByteBuffersDataOutput; import org.apache.lucene.store.ByteBuffersIndexOutput; import org.apache.lucene.util.Version; +import org.junit.Test; import org.opensearch.OpenSearchCorruptionException; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.settings.Settings; @@ -23,6 +24,7 @@ import org.opensearch.index.shard.IndexShardTestCase; import org.opensearch.index.store.StoreFileMetadata; import org.opensearch.indices.replication.checkpoint.MergedSegmentCheckpoint; +import org.opensearch.indices.replication.checkpoint.RemoteStoreMergedSegmentCheckpoint; import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint; import org.opensearch.indices.replication.common.ReplicationFailedException; import org.opensearch.indices.replication.common.ReplicationType; @@ -41,7 +43,8 @@ public class MergedSegmentReplicationTargetTests extends IndexShardTestCase { private MergedSegmentReplicationTarget mergedSegmentReplicationTarget; private IndexShard indexShard, spyIndexShard; - private MergedSegmentCheckpoint mergedSegment; + private MergedSegmentCheckpoint mergedSegmentCheckpoint; + private RemoteStoreMergedSegmentCheckpoint remoteStoreMergedSegmentCheckpoint; private ByteBuffersDataOutput buffer; private static final String SEGMENT_NAME = "_0.si"; @@ -71,7 +74,7 @@ public void setUp() throws Exception { try (ByteBuffersIndexOutput indexOutput = new ByteBuffersIndexOutput(buffer, "", null)) { testSegmentInfos.write(indexOutput); } - mergedSegment = new MergedSegmentCheckpoint( + mergedSegmentCheckpoint = new MergedSegmentCheckpoint( spyIndexShard.shardId(), spyIndexShard.getPendingPrimaryTerm(), 1, @@ -79,9 +82,18 @@ public void setUp() throws Exception { SI_SNAPSHOT, IndexFileNames.parseSegmentName(SEGMENT_NAME) ); + remoteStoreMergedSegmentCheckpoint = new RemoteStoreMergedSegmentCheckpoint( + spyIndexShard.shardId(), + spyIndexShard.getPendingPrimaryTerm(), + 1, + indexShard.getLatestReplicationCheckpoint().getCodec(), + SI_SNAPSHOT, + IndexFileNames.parseSegmentName(SEGMENT_NAME), + null + ); } - public void testSuccessfulResponse_startReplication() { + private void testSuccessfulResponse_startReplication(MergedSegmentCheckpoint checkpointMergedSegment) { SegmentReplicationSource segrepSource = new TestReplicationSource() { @Override @@ -119,7 +131,7 @@ public void getMergedSegmentFiles( SegmentReplicationTargetService.SegmentReplicationListener segRepListener = mock( SegmentReplicationTargetService.SegmentReplicationListener.class ); - mergedSegmentReplicationTarget = new MergedSegmentReplicationTarget(spyIndexShard, mergedSegment, segrepSource, segRepListener); + mergedSegmentReplicationTarget = new MergedSegmentReplicationTarget(spyIndexShard, checkpointMergedSegment, segrepSource, segRepListener); mergedSegmentReplicationTarget.startReplication(new ActionListener() { @Override @@ -133,12 +145,12 @@ public void onFailure(Exception e) { Assert.fail(); } }, (ReplicationCheckpoint checkpoint, IndexShard indexShard) -> { - assertEquals(mergedSegment, checkpoint); + assertEquals(mergedSegmentCheckpoint, checkpoint); assertEquals(indexShard, spyIndexShard); }); } - public void testFailureResponse_getMergedSegmentFiles() { + private void testFailureResponse_getMergedSegmentFiles(MergedSegmentCheckpoint checkpointMergedSegment) { Exception exception = new Exception("dummy failure"); SegmentReplicationSource segrepSource = new TestReplicationSource() { @@ -174,7 +186,7 @@ public void getMergedSegmentFiles( SegmentReplicationTargetService.SegmentReplicationListener segRepListener = mock( SegmentReplicationTargetService.SegmentReplicationListener.class ); - mergedSegmentReplicationTarget = new MergedSegmentReplicationTarget(spyIndexShard, mergedSegment, segrepSource, segRepListener); + mergedSegmentReplicationTarget = new MergedSegmentReplicationTarget(spyIndexShard, checkpointMergedSegment, segrepSource, segRepListener); mergedSegmentReplicationTarget.startReplication(new ActionListener() { @Override @@ -190,7 +202,7 @@ public void onFailure(Exception e) { }, mock(BiConsumer.class)); } - public void testFailure_differentSegmentFiles() throws IOException { + private void testFailure_differentSegmentFiles(MergedSegmentCheckpoint checkpointMergedSegment) throws IOException { SegmentReplicationSource segrepSource = new TestReplicationSource() { @Override @@ -225,7 +237,7 @@ public void getMergedSegmentFiles( SegmentReplicationTargetService.SegmentReplicationListener segRepListener = mock( SegmentReplicationTargetService.SegmentReplicationListener.class ); - mergedSegmentReplicationTarget = new MergedSegmentReplicationTarget(spyIndexShard, mergedSegment, segrepSource, segRepListener); + mergedSegmentReplicationTarget = new MergedSegmentReplicationTarget(spyIndexShard, checkpointMergedSegment, segrepSource, segRepListener); when(spyIndexShard.getSegmentMetadataMap()).thenReturn(SI_SNAPSHOT_DIFFERENT); mergedSegmentReplicationTarget.startReplication(new ActionListener() { @Override @@ -242,6 +254,30 @@ public void onFailure(Exception e) { }, mock(BiConsumer.class)); } + public void testFailure_differentSegmentFiles_remoteStoreEnabled() throws IOException { + testFailure_differentSegmentFiles(remoteStoreMergedSegmentCheckpoint); + } + + public void testFailure_differentSegmentFiles() throws IOException { + testFailure_differentSegmentFiles(mergedSegmentCheckpoint); + } + + public void testFailureResponse_getMergedSegmentFiles_remoteStoreEnabled() { + testFailureResponse_getMergedSegmentFiles(remoteStoreMergedSegmentCheckpoint); + } + + public void testFailureResponse_getMergedSegmentFiles() { + testFailureResponse_getMergedSegmentFiles(mergedSegmentCheckpoint); + } + + public void testSuccessfulResponse_startReplication_startReplication(){ + testSuccessfulResponse_startReplication(remoteStoreMergedSegmentCheckpoint); + } + + public void testSuccessfulResponse_startReplication() { + testSuccessfulResponse_startReplication(mergedSegmentCheckpoint); + } + @Override public void tearDown() throws Exception { super.tearDown(); diff --git a/server/src/test/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentActionTests.java b/server/src/test/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentActionTests.java new file mode 100644 index 0000000000000..0733ae81560cc --- /dev/null +++ b/server/src/test/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentActionTests.java @@ -0,0 +1,317 @@ +/* + * 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.indices.replication.checkpoint; + +import org.apache.lucene.codecs.Codec; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.ActionTestUtils; +import org.opensearch.action.support.PlainActionFuture; +import org.opensearch.action.support.replication.ReplicationMode; +import org.opensearch.action.support.replication.TransportReplicationAction; +import org.opensearch.cluster.action.shard.ShardStateAction; +import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.cluster.routing.AllocationId; +import org.opensearch.cluster.routing.ShardRouting; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.util.io.IOUtils; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.index.Index; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.IndexService; +import org.opensearch.index.shard.IndexShard; +import org.opensearch.index.store.Store; +import org.opensearch.indices.IndicesService; +import org.opensearch.indices.recovery.RecoverySettings; +import org.opensearch.indices.recovery.RecoveryState; +import org.opensearch.indices.replication.SegmentReplicationTargetService; +import org.opensearch.telemetry.tracing.noop.NoopTracer; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.transport.CapturingTransport; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import java.util.Collections; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.opensearch.index.remote.RemoteStoreTestsHelper.createIndexSettings; +import static org.opensearch.test.ClusterServiceUtils.createClusterService; +import static org.hamcrest.Matchers.sameInstance; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class RemoteStorePublishMergedSegmentActionTests extends OpenSearchTestCase { + + private ThreadPool threadPool; + private CapturingTransport transport; + private ClusterService clusterService; + private TransportService transportService; + private ShardStateAction shardStateAction; + + @Override + public void setUp() throws Exception { + super.setUp(); + threadPool = new TestThreadPool(getClass().getName()); + transport = new CapturingTransport(); + clusterService = createClusterService(threadPool); + transportService = transport.createTransportService( + clusterService.getSettings(), + threadPool, + TransportService.NOOP_TRANSPORT_INTERCEPTOR, + boundAddress -> clusterService.localNode(), + null, + Collections.emptySet(), + NoopTracer.INSTANCE + ); + transportService.start(); + transportService.acceptIncomingRequests(); + shardStateAction = new ShardStateAction(clusterService, transportService, null, null, threadPool); + } + + @Override + public void tearDown() throws Exception { + try { + IOUtils.close(transportService, clusterService, transport); + } finally { + terminate(threadPool); + } + super.tearDown(); + } + + public void testPublishMergedSegment() { + final IndicesService indicesService = mock(IndicesService.class); + + final Index index = new Index("index", "uuid"); + final IndexService indexService = mock(IndexService.class); + when(indicesService.indexServiceSafe(index)).thenReturn(indexService); + + final int id = randomIntBetween(0, 4); + final IndexShard indexShard = mock(IndexShard.class); + when(indexService.getShard(id)).thenReturn(indexShard); + + final ShardId shardId = new ShardId(index, id); + when(indexShard.shardId()).thenReturn(shardId); + + ShardRouting shardRouting = mock(ShardRouting.class); + AllocationId allocationId = mock(AllocationId.class); + RecoveryState recoveryState = mock(RecoveryState.class); + RecoverySettings recoverySettings = mock(RecoverySettings.class); + when(recoverySettings.getMergedSegmentReplicationTimeout()).thenReturn(new TimeValue(1000)); + when(shardRouting.allocationId()).thenReturn(allocationId); + when(allocationId.getId()).thenReturn("1"); + when(recoveryState.getTargetNode()).thenReturn(clusterService.localNode()); + when(indexShard.routingEntry()).thenReturn(shardRouting); + when(indexShard.getPendingPrimaryTerm()).thenReturn(1L); + when(indexShard.recoveryState()).thenReturn(recoveryState); + when(indexShard.getRecoverySettings()).thenReturn(recoverySettings); + when(indexShard.store()).thenReturn(mock(Store.class)); + + final SegmentReplicationTargetService mockTargetService = mock(SegmentReplicationTargetService.class); + + final RemoteStorePublishMergedSegmentAction action = new RemoteStorePublishMergedSegmentAction( + Settings.EMPTY, + transportService, + clusterService, + indicesService, + threadPool, + shardStateAction, + new ActionFilters(Collections.emptySet()), + mockTargetService + ); + + final MergedSegmentCheckpoint checkpoint = new RemoteStoreMergedSegmentCheckpoint( + indexShard.shardId(), + 1, + 1111, + Codec.getDefault().getName(), + Collections.emptyMap(), + "_1", + null + ); + + action.publish(indexShard, checkpoint); + } + + public void testPublishMergedSegmentActionOnPrimary() { + final IndicesService indicesService = mock(IndicesService.class); + + final Index index = new Index("index", "uuid"); + final IndexService indexService = mock(IndexService.class); + when(indicesService.indexServiceSafe(index)).thenReturn(indexService); + + final int id = randomIntBetween(0, 4); + final IndexShard indexShard = mock(IndexShard.class); + when(indexService.getShard(id)).thenReturn(indexShard); + + final ShardId shardId = new ShardId(index, id); + when(indexShard.shardId()).thenReturn(shardId); + + final SegmentReplicationTargetService mockTargetService = mock(SegmentReplicationTargetService.class); + + final RemoteStorePublishMergedSegmentAction action = new RemoteStorePublishMergedSegmentAction( + Settings.EMPTY, + transportService, + clusterService, + indicesService, + threadPool, + shardStateAction, + new ActionFilters(Collections.emptySet()), + mockTargetService + ); + + final RemoteStoreMergedSegmentCheckpoint checkpoint = new RemoteStoreMergedSegmentCheckpoint( + indexShard.shardId(), + 1, + 1111, + Codec.getDefault().getName(), + Collections.emptyMap(), + "_1", + null + ); + final RemoteStorePublishMergedSegmentRequest request = new RemoteStorePublishMergedSegmentRequest(checkpoint); + + action.shardOperationOnPrimary(request, indexShard, ActionTestUtils.assertNoFailureListener(result -> { + // we should forward the request containing the current publish checkpoint to the replica + assertThat(result.replicaRequest(), sameInstance(request)); + })); + } + + public void testPublishMergedSegmentActionOnReplica() { + final IndicesService indicesService = mock(IndicesService.class); + + final Index index = new Index("index", "uuid"); + final IndexService indexService = mock(IndexService.class); + when(indicesService.indexServiceSafe(index)).thenReturn(indexService); + final int id = randomIntBetween(0, 4); + final IndexShard indexShard = mock(IndexShard.class); + when(indexService.getShard(id)).thenReturn(indexShard); + + final ShardId shardId = new ShardId(index, id); + when(indexShard.shardId()).thenReturn(shardId); + when(indexShard.indexSettings()).thenReturn( + createIndexSettings(false, Settings.builder().put(IndexMetadata.INDEX_REPLICATION_TYPE_SETTING.getKey(), "SEGMENT").build()) + ); + final SegmentReplicationTargetService mockTargetService = mock(SegmentReplicationTargetService.class); + + final RemoteStorePublishMergedSegmentAction action = new RemoteStorePublishMergedSegmentAction( + Settings.EMPTY, + transportService, + clusterService, + indicesService, + threadPool, + shardStateAction, + new ActionFilters(Collections.emptySet()), + mockTargetService + ); + + final RemoteStoreMergedSegmentCheckpoint checkpoint = createCheckpoint(indexShard); + + final RemoteStorePublishMergedSegmentRequest request = new RemoteStorePublishMergedSegmentRequest(checkpoint); + + final PlainActionFuture listener = PlainActionFuture.newFuture(); + action.shardOperationOnReplica(request, indexShard, listener); + final TransportReplicationAction.ReplicaResult result = listener.actionGet(); + + // onNewMergedSegmentCheckpoint should be called on shard with checkpoint request + verify(mockTargetService, times(1)).onNewMergedSegmentCheckpoint(checkpoint, indexShard); + + // the result should indicate success + final AtomicBoolean success = new AtomicBoolean(); + result.runPostReplicaActions(ActionListener.wrap(r -> success.set(true), e -> fail(e.toString()))); + assertTrue(success.get()); + + } + + public void testPublishMergedSegmentActionOnDocrepReplicaDuringMigration() { + final IndicesService indicesService = mock(IndicesService.class); + + final Index index = new Index("index", "uuid"); + final IndexService indexService = mock(IndexService.class); + when(indicesService.indexServiceSafe(index)).thenReturn(indexService); + final int id = randomIntBetween(0, 4); + final IndexShard indexShard = mock(IndexShard.class); + when(indexService.getShard(id)).thenReturn(indexShard); + + final ShardId shardId = new ShardId(index, id); + when(indexShard.shardId()).thenReturn(shardId); + when(indexShard.indexSettings()).thenReturn(createIndexSettings(false)); + final SegmentReplicationTargetService mockTargetService = mock(SegmentReplicationTargetService.class); + + final RemoteStorePublishMergedSegmentAction action = new RemoteStorePublishMergedSegmentAction( + Settings.EMPTY, + transportService, + clusterService, + indicesService, + threadPool, + shardStateAction, + new ActionFilters(Collections.emptySet()), + mockTargetService + ); + + final RemoteStoreMergedSegmentCheckpoint checkpoint = createCheckpoint(indexShard); + + final RemoteStorePublishMergedSegmentRequest request = new RemoteStorePublishMergedSegmentRequest(checkpoint); + + final PlainActionFuture listener = PlainActionFuture.newFuture(); + action.shardOperationOnReplica(request, indexShard, listener); + final TransportReplicationAction.ReplicaResult result = listener.actionGet(); + // no interaction with SegmentReplicationTargetService object + verify(mockTargetService, never()).onNewMergedSegmentCheckpoint(any(), any()); + // the result should indicate success + final AtomicBoolean success = new AtomicBoolean(); + result.runPostReplicaActions(ActionListener.wrap(r -> success.set(true), e -> fail(e.toString()))); + assertTrue(success.get()); + } + + public void testGetReplicationModeWithRemoteTranslog() { + final RemoteStorePublishMergedSegmentAction action = createAction(); + final IndexShard indexShard = mock(IndexShard.class); + when(indexShard.indexSettings()).thenReturn(createIndexSettings(true)); + assertEquals(ReplicationMode.FULL_REPLICATION, action.getReplicationMode(indexShard)); + } + + public void testGetReplicationModeWithLocalTranslog() { + final RemoteStorePublishMergedSegmentAction action = createAction(); + final IndexShard indexShard = mock(IndexShard.class); + when(indexShard.indexSettings()).thenReturn(createIndexSettings(false)); + assertEquals(ReplicationMode.FULL_REPLICATION, action.getReplicationMode(indexShard)); + } + + private RemoteStorePublishMergedSegmentAction createAction() { + return new RemoteStorePublishMergedSegmentAction( + Settings.EMPTY, + transportService, + clusterService, + mock(IndicesService.class), + threadPool, + shardStateAction, + new ActionFilters(Collections.emptySet()), + mock(SegmentReplicationTargetService.class) + ); + } + + private RemoteStoreMergedSegmentCheckpoint createCheckpoint(IndexShard indexShard) { + return new RemoteStoreMergedSegmentCheckpoint( + indexShard.shardId(), + 1, + 1111, + Codec.getDefault().getName(), + Collections.emptyMap(), + "_1", + null + ); + } + +} diff --git a/server/src/test/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentRequestTests.java b/server/src/test/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentRequestTests.java new file mode 100644 index 0000000000000..55ab4281d1488 --- /dev/null +++ b/server/src/test/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentRequestTests.java @@ -0,0 +1,67 @@ +/* + * 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.indices.replication.checkpoint; + +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.index.Index; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.Locale; +import java.util.Map; +import java.util.Objects; + +import static org.hamcrest.core.IsEqual.equalTo; + +public class RemoteStorePublishMergedSegmentRequestTests extends OpenSearchTestCase { + + public void testPublishMergedSegmentRequest() { + RemoteStoreMergedSegmentCheckpoint checkpoint = new RemoteStoreMergedSegmentCheckpoint( + new ShardId(new Index("1", "1"), 0), + 0, + 0, + "", + Collections.emptyMap(), + "_0", + null + ); + + RemoteStorePublishMergedSegmentRequest request = new RemoteStorePublishMergedSegmentRequest(checkpoint); + assertNull(request.validate()); + assertEquals(checkpoint, request.getMergedSegment()); + assertEquals(Objects.hash(checkpoint), request.hashCode()); + assertEquals(checkpoint.getLocalToRemoteSegmentFilenameMap(), Map.of()); + checkpoint.updateLocalToRemoteSegmentFilenameMap("_0", "_0__uuid"); + assertEquals(checkpoint.getLocalToRemoteSegmentFilenameMap(), Map.of("_0", "_0__uuid")); + } + + public void testSerialize() throws Exception { + RemoteStoreMergedSegmentCheckpoint checkpoint = new RemoteStoreMergedSegmentCheckpoint( + new ShardId(new Index("1", "1"), 0), + 0, + 0, + "", + Collections.emptyMap(), + "_0", + null + ); + RemoteStorePublishMergedSegmentRequest originalRequest = new RemoteStorePublishMergedSegmentRequest(checkpoint); + RemoteStorePublishMergedSegmentRequest cloneRequest; + try (BytesStreamOutput out = new BytesStreamOutput()) { + originalRequest.writeTo(out); + try (StreamInput in = out.bytes().streamInput()) { + cloneRequest = new RemoteStorePublishMergedSegmentRequest(in); + } + } + assertThat(cloneRequest, equalTo(originalRequest)); + assertThat(cloneRequest.getMergedSegment(), equalTo(originalRequest.getMergedSegment())); + } +} diff --git a/server/src/test/java/org/opensearch/repositories/RepositoriesServiceTests.java b/server/src/test/java/org/opensearch/repositories/RepositoriesServiceTests.java index 1ec6d320762f2..0938242ffc918 100644 --- a/server/src/test/java/org/opensearch/repositories/RepositoriesServiceTests.java +++ b/server/src/test/java/org/opensearch/repositories/RepositoriesServiceTests.java @@ -724,6 +724,11 @@ public long getRemoteDownloadThrottleTimeInNanos() { return 0; } + @Override + public long getLowPriorityRemoteDownloadThrottleTimeInNanos() { + return 0; + } + @Override public String startVerification() { return null; diff --git a/test/framework/src/main/java/org/opensearch/index/shard/RestoreOnlyRepository.java b/test/framework/src/main/java/org/opensearch/index/shard/RestoreOnlyRepository.java index 846f54bd13b40..fe549fdbb992e 100644 --- a/test/framework/src/main/java/org/opensearch/index/shard/RestoreOnlyRepository.java +++ b/test/framework/src/main/java/org/opensearch/index/shard/RestoreOnlyRepository.java @@ -175,6 +175,12 @@ public long getRemoteDownloadThrottleTimeInNanos() { return 0; } + @Override + public long getLowPriorityRemoteDownloadThrottleTimeInNanos() { + return 0; + } + + @Override public String startVerification() { return null;