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 315360fa4d99e..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.MergeSegmentCheckpoint; -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; @@ -1864,19 +1861,30 @@ 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 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); - return new MergeSegmentCheckpoint( + 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(), segmentMetadataMap.values().stream().mapToLong(StoreFileMetadata::length).sum(), 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/shard/RemoteStoreUploader.java b/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploader.java new file mode 100644 index 0000000000000..39ce83215f43e --- /dev/null +++ b/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploader.java @@ -0,0 +1,29 @@ +/* + * 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.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 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 new file mode 100644 index 0000000000000..b0e06ca978222 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/shard/RemoteStoreUploaderService.java @@ -0,0 +1,104 @@ +/* + * 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.index.CorruptIndexException; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FilterDirectory; +import org.apache.lucene.store.IOContext; +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.store.CompositeDirectory; +import org.opensearch.index.store.RemoteSegmentStoreDirectory; + +import java.util.Collection; +import java.util.Map; +import java.util.function.Function; + +/** + * 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; + + private final IndexShard indexShard; + private final Directory storeDirectory; + private final RemoteSegmentStoreDirectory remoteDirectory; + + public RemoteStoreUploaderService(IndexShard indexShard, Directory storeDirectory, RemoteSegmentStoreDirectory remoteDirectory) { + logger = Loggers.getLogger(getClass(), indexShard.shardId()); + this.indexShard = indexShard; + this.storeDirectory = storeDirectory; + this.remoteDirectory = remoteDirectory; + } + + @Override + public void uploadSegments( + Collection localSegments, + Map localSegmentsSizeMap, + ActionListener listener, + Function, UploadListener> uploadListenerFunction + ) { + if (localSegments.isEmpty()) { + logger.debug("No new segments to upload in uploadNewSegments"); + listener.onResponse(null); + return; + } + + logger.debug("Effective new segments files to upload {}", localSegments); + ActionListener> mappedListener = ActionListener.map(listener, resp -> null); + GroupedActionListener batchUploadListener = new GroupedActionListener<>(mappedListener, localSegments.size()); + Directory directory = ((FilterDirectory) (((FilterDirectory) storeDirectory).getDelegate())).getDelegate(); + + for (String localSegment : localSegments) { + // Initializing listener here to ensure that the stats increment operations are thread-safe + UploadListener statsListener = uploadListenerFunction.apply(localSegmentsSizeMap); + ActionListener aggregatedListener = ActionListener.wrap(resp -> { + 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(localSegment); + } + }, ex -> { + logger.warn(() -> new ParameterizedMessage("Exception: [{}] while uploading segment files", ex), ex); + if (ex instanceof CorruptIndexException) { + indexShard.failShard(ex.getMessage(), ex); + } + statsListener.onFailure(localSegment); + batchUploadListener.onFailure(ex); + }); + statsListener.beforeUpload(localSegment); + // Place where the actual upload is happening + remoteDirectory.copyFrom(storeDirectory, localSegment, 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()); + } +} 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..ded50f355e901 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,52 @@ 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) { + 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: Do we have a better way to check this? + return name.contains(SEGMENT_NAME_UUID_SEPARATOR); } /** @@ -988,6 +1045,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..899340026b7a6 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) {} + public RemoteStoreFileDownloader(ShardId shardId, ThreadPool threadPool, RecoverySettings recoverySettings) { this.logger = Loggers.getLogger(RemoteStoreFileDownloader.class, shardId); this.threadPool = threadPool; @@ -72,16 +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( + public void download( Directory source, Directory destination, Directory secondDestination, - Collection toDownloadSegments, + Collection toDownloadSegments, Runnable onFileCompletion ) throws InterruptedException, IOException { final CancellableThreads cancellableThreads = new CancellableThreads(); @@ -105,16 +109,16 @@ public void download( } } - private void downloadInternal( + 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 +134,36 @@ 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(String localFilename, String remoteFilename)) { + fileSrc = remoteFilename; + fileDest = 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/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..2e3d2e0868d78 --- /dev/null +++ b/server/src/main/java/org/opensearch/indices/replication/ActiveMergesSegmentRegistry.java @@ -0,0 +1,115 @@ +/* + * 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.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; + +/** + * 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 segmentMetadataMap = new ConcurrentHashMap<>(); + final Set filenameRegistry = ConcurrentHashMap.newKeySet(); // package-private for tests + private final ReentrantLock lock = new ReentrantLock(); + + private static class HOLDER { + private static final ActiveMergesSegmentRegistry INSTANCE = new ActiveMergesSegmentRegistry(); + } + + private ActiveMergesSegmentRegistry() {}; + + public static ActiveMergesSegmentRegistry getInstance() { + return HOLDER.INSTANCE; + } + + /** + * 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."); + } + filenameRegistry.add(localSegmentFilename); + } finally { + lock.unlock(); + } + } + + /** + * 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"); + } + segmentMetadataMap.put(localSegmentFilename, metadata); + filenameRegistry.add(metadata.getUploadedFilename()); + } finally { + lock.unlock(); + } + } + + /** + * 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 filenameRegistry.contains(segmentFilename); + } + + 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(); + } + + public boolean canDelete(@NonNull String segmentFilename) { + return contains(segmentFilename) == false; + } + + public Map segmentMetadataMap() { + return Collections.unmodifiableMap(segmentMetadataMap); + } + + 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..88389338a0bf4 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) + ) + ); + } + + indexShard.getFileDownloader() + .download( + 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..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.debug( + 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.debug( + logger.trace( "[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/MergeSegmentCheckpoint.java b/server/src/main/java/org/opensearch/indices/replication/checkpoint/MergedSegmentCheckpoint.java similarity index 92% 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..8bc1b596100ea 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 { - private final String segmentName; +public class MergedSegmentCheckpoint extends ReplicationCheckpoint { + 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..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 @@ -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, ReplicationCheckpoint 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..145d96e1aee6c --- /dev/null +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishMergedSegmentActionProvider.java @@ -0,0 +1,69 @@ +/* + * 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.common.util.FeatureFlags; +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() { + if (FeatureFlags.isEnabled(FeatureFlags.MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING) == false) { + return null; + } + // TODO: FIX THIS + if (false) {// || clusterService.localNode().isRemoteStoreNode() == 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..dd430022e71ce --- /dev/null +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStoreMergedSegmentCheckpoint.java @@ -0,0 +1,91 @@ +/* + * 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.Nullable; +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.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +/** + * Represents a remote store merged segment checkpoint. + * + * @opensearch.internal + */ +@ExperimentalApi +public class RemoteStoreMergedSegmentCheckpoint extends MergedSegmentCheckpoint { + private final Map localToRemoteSegmentFilenameMap; + + public RemoteStoreMergedSegmentCheckpoint( + ShardId shardId, + long primaryTerm, + long length, + String codec, + Map metadataMap, + String segmentName, + @Nullable Map localToRemoteSegmentFilenameMap + ) { + super(shardId, primaryTerm, length, codec, metadataMap, segmentName); + this.localToRemoteSegmentFilenameMap = localToRemoteSegmentFilenameMap == null ? new HashMap<>() : localToRemoteSegmentFilenameMap; + } + + public RemoteStoreMergedSegmentCheckpoint(StreamInput in) throws IOException { + super(in); + this.localToRemoteSegmentFilenameMap = in.readMap(StreamInput::readString, StreamInput::readString); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeMap(getLocalToRemoteSegmentFilenameMap(), StreamOutput::writeString, StreamOutput::writeString); + } + + @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) + && Objects.equals(getShardId(), that.getShardId()) + && getCodec().equals(that.getCodec()); + } + + @Override + public String toString() { + return "RemoteStoreMergedSegmentCheckpoint{" + + "shardId=" + + getShardId() + + ", primaryTerm=" + + getPrimaryTerm() + + ", segmentName=" + + getSegmentName() + + ", localToRemoteSegmentFilenameSize=" + + getLocalToRemoteSegmentFilenameMap().size() + + '}'; + } + + public Map getLocalToRemoteSegmentFilenameMap() { + return this.localToRemoteSegmentFilenameMap; + } + + 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 new file mode 100644 index 0000000000000..e9fc61c6470fb --- /dev/null +++ b/server/src/main/java/org/opensearch/indices/replication/checkpoint/RemoteStorePublishMergedSegmentAction.java @@ -0,0 +1,179 @@ +/* + * 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.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; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +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; + +public class RemoteStorePublishMergedSegmentAction extends AbstractPublishCheckpointAction implements MergedSegmentPublisher.PublishAction { + + public static final String ACTION_NAME = "indices:admin/remote_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) { + 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) { + if(! (checkpoint instanceof RemoteStoreMergedSegmentCheckpoint mergedSegmentCheckpoint)) { + throw new AssertionError("Expected checkpoint to be an instance of " + RemoteStoreMergedSegmentCheckpoint.class); + } + + publishMergedSegmentsToRemoteStore(indexShard, mergedSegmentCheckpoint); + doPublish(indexShard, + checkpoint, + new RemoteStorePublishMergedSegmentRequest(mergedSegmentCheckpoint), + "segrep_publish_merged_segment", + true, + indexShard.getRecoverySettings().getMergedSegmentReplicationTimeout() + ); + } + + private void publishMergedSegmentsToRemoteStore(IndexShard indexShard, RemoteStoreMergedSegmentCheckpoint checkpoint) { + RemoteStoreUploaderService remoteStoreUploaderService = getRemoteStoreUploaderService(indexShard); + Collection segmentsToUpload = checkpoint.getMetadataMap().keySet(); + + Map segmentsSizeMap = checkpoint + .getMetadataMap() + .entrySet() + .stream() + .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().length())); + + final CountDownLatch latch = new CountDownLatch(segmentsToUpload.size()); + + // TODO: Upload in low priority + remoteStoreUploaderService.uploadSegments( + segmentsToUpload, + segmentsSizeMap, + new ActionListener() { + @Override + public void onResponse(Void unused) { + if(logger.isTraceEnabled() == true) { + logger.trace("Successfully uploaded segments {} to remote store", segmentsToUpload); + } + } + + @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) { + activeMergesSegmentRegistry.register(file); + } + + @Override + public void onSuccess(String file) { + checkpoint.updateLocalToRemoteSegmentFilenameMap( + file, + activeMergesSegmentRegistry.getExistingRemoteSegmentFilename(file) + ); + latch.countDown(); + } + + @Override + public void onFailure(String file) { + segmentsToUpload.forEach(activeMergesSegmentRegistry::unregister); + /** + * TODO: abort merge + */ + } + } + ); + try { + 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: abort merge properly here + } + } + + /** + * TODO: REBASE ONCE UPLOAD CHANGES ARE COMPLETE + */ + private RemoteStoreUploaderService getRemoteStoreUploaderService(IndexShard indexShard) { + return new RemoteStoreUploaderService( + indexShard, + indexShard.store().directory(), + indexShard.getRemoteDirectory() + ); + } +} 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/main/java/org/opensearch/repositories/FilterRepository.java b/server/src/main/java/org/opensearch/repositories/FilterRepository.java index 188d61e00f206..60b1ab9cd871a 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.getLowPriorityRemoteDownloadThrottleTimeInNanos(); + } + @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); } 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 53c8e29acf2c9..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; @@ -22,7 +23,8 @@ 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.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 MergeSegmentCheckpoint 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 MergeSegmentCheckpoint( + 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/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, 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;