Skip to content

Clone index input in multipart upload flow #19072

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Aug 19, 2025
Merged

Clone index input in multipart upload flow #19072

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
ee8717d
Clone index input in multipart upload flow
sachinpkale Aug 14, 2025
cf27be8
Fix spotless error
sachinpkale Aug 14, 2025
18138d7
Fix for segments_N file
sachinpkale Aug 14, 2025
0464761
Fix spotless error
sachinpkale Aug 14, 2025
b04e72a
Fix test
sachinpkale Aug 18, 2025
b1069f8
Make sure to close IndexInput on exception
sachinpkale Aug 19, 2025
93b9237
Fix spotless errors
sachinpkale Aug 19, 2025
c4cf7f4
Address PR comments
sachinpkale Aug 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 71 additions & 57 deletions server/src/main/java/org/opensearch/index/store/RemoteDirectory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -397,68 +397,82 @@ private void uploadBlob(
assert ioContext != IOContext.READONCE : "Remote upload will fail with IoContext.READONCE";
long expectedChecksum = calculateChecksumOfChecksum(from, src);
long contentLength;
try (IndexInput indexInput = from.openInput(src, ioContext)) {
IndexInput indexInput = from.openInput(src, ioContext);
try {
contentLength = indexInput.length();
}
boolean remoteIntegrityEnabled = false;
if (getBlobContainer() instanceof AsyncMultiStreamBlobContainer) {
remoteIntegrityEnabled = ((AsyncMultiStreamBlobContainer) getBlobContainer()).remoteIntegrityCheckSupported();
}
lowPriorityUpload = lowPriorityUpload || contentLength > ByteSizeUnit.GB.toBytes(15);
RemoteTransferContainer.OffsetRangeInputStreamSupplier offsetRangeInputStreamSupplier;
if (lowPriorityUpload) {
offsetRangeInputStreamSupplier = (size, position) -> lowPriorityUploadRateLimiter.apply(
new OffsetRangeIndexInputStream(from.openInput(src, ioContext), size, position)
);
} else {
offsetRangeInputStreamSupplier = (size, position) -> uploadRateLimiter.apply(
new OffsetRangeIndexInputStream(from.openInput(src, ioContext), size, position)
);
}
RemoteTransferContainer remoteTransferContainer = new RemoteTransferContainer(
src,
remoteFileName,
contentLength,
true,
lowPriorityUpload ? WritePriority.LOW : WritePriority.NORMAL,
offsetRangeInputStreamSupplier,
expectedChecksum,
remoteIntegrityEnabled
);
ActionListener<Void> completionListener = ActionListener.wrap(resp -> {
try {
postUploadRunner.run();
listener.onResponse(null);
} catch (Exception e) {
logger.error(() -> new ParameterizedMessage("Exception in segment postUpload for file [{}]", src), e);
listener.onFailure(e);
}
}, ex -> {
logger.error(() -> new ParameterizedMessage("Failed to upload blob {}", src), ex);
IOException corruptIndexException = ExceptionsHelper.unwrapCorruption(ex);
if (corruptIndexException != null) {
listener.onFailure(corruptIndexException);
return;
}
Throwable throwable = ExceptionsHelper.unwrap(ex, CorruptFileException.class);
if (throwable != null) {
CorruptFileException corruptFileException = (CorruptFileException) throwable;
listener.onFailure(new CorruptIndexException(corruptFileException.getMessage(), corruptFileException.getFileName()));
return;
boolean remoteIntegrityEnabled = false;
if (getBlobContainer() instanceof AsyncMultiStreamBlobContainer) {
remoteIntegrityEnabled = ((AsyncMultiStreamBlobContainer) getBlobContainer()).remoteIntegrityCheckSupported();
}
listener.onFailure(ex);
});
lowPriorityUpload = lowPriorityUpload || contentLength > ByteSizeUnit.GB.toBytes(15);
RemoteTransferContainer.OffsetRangeInputStreamSupplier offsetRangeInputStreamSupplier;

completionListener = ActionListener.runBefore(completionListener, () -> {
try {
remoteTransferContainer.close();
} catch (Exception e) {
logger.warn("Error occurred while closing streams", e);
if (lowPriorityUpload) {
offsetRangeInputStreamSupplier = (size, position) -> lowPriorityUploadRateLimiter.apply(
new OffsetRangeIndexInputStream(indexInput.clone(), size, position)
);
} else {
offsetRangeInputStreamSupplier = (size, position) -> uploadRateLimiter.apply(
new OffsetRangeIndexInputStream(indexInput.clone(), size, position)
);
}
});
RemoteTransferContainer remoteTransferContainer = new RemoteTransferContainer(
src,
remoteFileName,
contentLength,
true,
lowPriorityUpload ? WritePriority.LOW : WritePriority.NORMAL,
offsetRangeInputStreamSupplier,
expectedChecksum,
remoteIntegrityEnabled
);
ActionListener<Void> completionListener = ActionListener.wrap(resp -> {
try {
postUploadRunner.run();
listener.onResponse(null);
} catch (Exception e) {
logger.error(() -> new ParameterizedMessage("Exception in segment postUpload for file [{}]", src), e);
listener.onFailure(e);
}
}, ex -> {
logger.error(() -> new ParameterizedMessage("Failed to upload blob {}", src), ex);
IOException corruptIndexException = ExceptionsHelper.unwrapCorruption(ex);
if (corruptIndexException != null) {
listener.onFailure(corruptIndexException);
return;
}
Throwable throwable = ExceptionsHelper.unwrap(ex, CorruptFileException.class);
if (throwable != null) {
CorruptFileException corruptFileException = (CorruptFileException) throwable;
listener.onFailure(new CorruptIndexException(corruptFileException.getMessage(), corruptFileException.getFileName()));
return;
}
listener.onFailure(ex);
});

WriteContext writeContext = remoteTransferContainer.createWriteContext();
((AsyncMultiStreamBlobContainer) blobContainer).asyncBlobUpload(writeContext, completionListener);
completionListener = ActionListener.runBefore(completionListener, () -> {
try {
remoteTransferContainer.close();
} catch (Exception e) {
logger.warn("Error occurred while closing streams", e);
}
});

completionListener = ActionListener.runAfter(completionListener, () -> {
try {
indexInput.close();
Copy link
Member

@ashking94 ashking94 Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good to add a debug/trace log line for places where the close is happening to check in future if any potential leak happens.

} catch (IOException e) {
logger.warn("Error occurred while closing index input", e);
}
});

WriteContext writeContext = remoteTransferContainer.createWriteContext();
((AsyncMultiStreamBlobContainer) blobContainer).asyncBlobUpload(writeContext, completionListener);
} catch (Exception e) {
logger.warn("Exception while calling asyncBlobUpload, closing IndexInput to avoid leak");
indexInput.close();
throw e;
}
}

private long calculateChecksumOfChecksum(Directory directory, String file) throws IOException {
Expand Down
18 changes: 11 additions & 7 deletions server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import org.apache.logging.log4j.Logger;
import org.apache.logging.log4j.message.ParameterizedMessage;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.index.IndexFileNames;
import org.apache.lucene.index.SegmentCommitInfo;
import org.apache.lucene.index.SegmentInfo;
import org.apache.lucene.index.SegmentInfos;
Expand Down Expand Up @@ -599,13 +600,16 @@ public IndexInput openBlockInput(String name, long position, long length, IOCont
public void copyFrom(Directory from, String src, IOContext context, ActionListener<Void> listener, boolean lowPriorityUpload) {
try {
final String remoteFileName = getNewRemoteSegmentFilename(src);
boolean uploaded = remoteDataDirectory.copyFrom(from, src, remoteFileName, context, () -> {
try {
postUpload(from, src, remoteFileName, getChecksumOfLocalFile(from, src));
} catch (IOException e) {
throw new RuntimeException("Exception in segment postUpload for file " + src, e);
}
}, listener, lowPriorityUpload);
boolean uploaded = false;
if (src.startsWith(IndexFileNames.SEGMENTS) == false) {
uploaded = remoteDataDirectory.copyFrom(from, src, remoteFileName, context, () -> {
try {
postUpload(from, src, remoteFileName, getChecksumOfLocalFile(from, src));
} catch (IOException e) {
throw new RuntimeException("Exception in segment postUpload for file " + src, e);
}
}, listener, lowPriorityUpload);
}
if (uploaded == false) {
copyFrom(from, src, src, context);
listener.onResponse(null);
Expand Down
10 changes: 8 additions & 2 deletions server/src/test/java/org/opensearch/index/shard/RemoteStoreRefreshListenerTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,8 @@
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import org.mockito.stubbing.Answer;

import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_REPLICATION_TYPE;
import static org.opensearch.index.store.RemoteSegmentStoreDirectory.METADATA_FILES_TO_FETCH;
import static org.opensearch.test.RemoteStoreTestUtils.createMetadataFileBytes;
Expand Down Expand Up @@ -822,7 +824,7 @@ private Tuple<RemoteStoreRefreshListener, RemoteStoreStatsTrackerFactory> mockIn
when(shard.getRemoteStoreSettings()).thenReturn(remoteStoreSettings);
if (testUploadTimeout) {
when(remoteStoreSettings.getClusterRemoteSegmentTransferTimeout()).thenReturn(TimeValue.timeValueMillis(10));
doAnswer(invocation -> {
Answer<Boolean> answer = invocation -> {
ActionListener<Void> actionListener = invocation.getArgument(5);
indexShard.getThreadPool().executor(ThreadPool.Names.GENERIC).execute(() -> {
try {
Expand All @@ -833,7 +835,11 @@ private Tuple<RemoteStoreRefreshListener, RemoteStoreStatsTrackerFactory> mockIn
actionListener.onResponse(null);
});
return true;
}).when(remoteDirectory).copyFrom(any(), any(), any(), any(), any(), any(ActionListener.class), any(Boolean.class));
};
doAnswer(answer).when(remoteDirectory)
.copyFrom(any(), any(), any(), any(), any(), any(ActionListener.class), any(Boolean.class));

doAnswer(answer).when(remoteDirectory).copyFrom(any(), any(), any(), any());
}

RemoteStoreRefreshListener refreshListener = new RemoteStoreRefreshListener(
Expand Down
Loading