-
Notifications
You must be signed in to change notification settings - Fork 1.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add async blob read and download support using multiple streams
Signed-off-by: Kunal Kotwani <[email protected]>
- Loading branch information
1 parent
a08d588
commit 62fbe0d
Showing
15 changed files
with
748 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
46 changes: 46 additions & 0 deletions
46
server/src/main/java/org/opensearch/common/blobstore/stream/read/ReadContext.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,46 @@ | ||
| /* | ||
| * 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.common.blobstore.stream.read; | ||
|
|
||
| import org.opensearch.common.io.InputStreamContainer; | ||
|
|
||
| import java.util.List; | ||
|
|
||
| /** | ||
| * ReadContext is used to encapsulate all data needed by <code>BlobContainer#readBlobAsync</code> | ||
| * | ||
| * @opensearch.internal | ||
| */ | ||
| public class ReadContext { | ||
| private final long blobSize; | ||
| private final List<InputStreamContainer> partStreams; | ||
| private final String blobChecksum; | ||
|
|
||
| public ReadContext(long blobSize, List<InputStreamContainer> partStreams, String blobChecksum) { | ||
| this.blobSize = blobSize; | ||
| this.partStreams = partStreams; | ||
| this.blobChecksum = blobChecksum; | ||
| } | ||
|
|
||
| public String getBlobChecksum() { | ||
| return blobChecksum; | ||
| } | ||
|
|
||
| public int getNumberOfParts() { | ||
| return partStreams.size(); | ||
| } | ||
|
|
||
| public long getBlobSize() { | ||
| return blobSize; | ||
| } | ||
|
|
||
| public List<InputStreamContainer> getPartStreams() { | ||
| return partStreams; | ||
| } | ||
| } |
49 changes: 49 additions & 0 deletions
49
...ain/java/org/opensearch/common/blobstore/stream/read/listener/FileCompletionListener.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,49 @@ | ||
| /* | ||
| * 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.common.blobstore.stream.read.listener; | ||
|
|
||
| import org.opensearch.core.action.ActionListener; | ||
|
|
||
| import java.util.Collections; | ||
| import java.util.HashSet; | ||
| import java.util.Set; | ||
|
|
||
| /** | ||
| * FileCompletionListener listens for completion of fetch on all the streams for a file, where | ||
| * individual streams are handled using {@link StreamCompletionListener}. The {@link StreamCompletionListener}(s) | ||
| * hold a reference to the file completion listener to be notified. | ||
| * | ||
| * @opensearch.internal | ||
| */ | ||
| public class FileCompletionListener implements ActionListener<Integer> { | ||
| private final int numberOfParts; | ||
| private final String fileName; | ||
| private final Set<Integer> completedParts; | ||
| private final ActionListener<String> completionListener; | ||
|
|
||
| public FileCompletionListener(int numberOfParts, String fileName, ActionListener<String> completionListener) { | ||
| this.completedParts = Collections.synchronizedSet(new HashSet<>()); | ||
| this.numberOfParts = numberOfParts; | ||
| this.fileName = fileName; | ||
| this.completionListener = completionListener; | ||
| } | ||
|
|
||
| @Override | ||
| public void onResponse(Integer partNumber) { | ||
| completedParts.add(partNumber); | ||
| if (completedParts.size() == numberOfParts) { | ||
| completionListener.onResponse(fileName); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public void onFailure(Exception e) { | ||
| completionListener.onFailure(e); | ||
| } | ||
| } |
65 changes: 65 additions & 0 deletions
65
...c/main/java/org/opensearch/common/blobstore/stream/read/listener/ReadContextListener.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,65 @@ | ||
| /* | ||
| * 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.common.blobstore.stream.read.listener; | ||
|
|
||
| import org.apache.logging.log4j.LogManager; | ||
| import org.apache.logging.log4j.Logger; | ||
| import org.opensearch.action.support.ThreadedActionListener; | ||
| import org.opensearch.common.blobstore.stream.read.ReadContext; | ||
| import org.opensearch.core.action.ActionListener; | ||
| import org.opensearch.threadpool.ThreadPool; | ||
|
|
||
| import java.nio.file.Path; | ||
| import java.util.concurrent.atomic.AtomicBoolean; | ||
|
|
||
| /** | ||
| * ReadContextListener orchestrates the async file fetch from the {@link org.opensearch.common.blobstore.BlobContainer} | ||
| * using a {@link ReadContext} callback. On response, it spawns off the download using multiple streams which are | ||
| * spread across a {@link ThreadPool} executor. | ||
| * | ||
| * @opensearch.internal | ||
| */ | ||
| public class ReadContextListener implements ActionListener<ReadContext> { | ||
|
|
||
| private final String fileName; | ||
| private final Path fileLocation; | ||
| private final ThreadPool threadPool; | ||
| private final ActionListener<String> completionListener; | ||
| private static final Logger logger = LogManager.getLogger(ReadContextListener.class); | ||
|
|
||
| public ReadContextListener(String fileName, Path fileLocation, ThreadPool threadPool, ActionListener<String> completionListener) { | ||
| this.fileName = fileName; | ||
| this.fileLocation = fileLocation; | ||
| this.threadPool = threadPool; | ||
| this.completionListener = completionListener; | ||
| } | ||
|
|
||
| @Override | ||
| public void onResponse(ReadContext readContext) { | ||
| final int numParts = readContext.getNumberOfParts(); | ||
| final AtomicBoolean anyPartStreamFailed = new AtomicBoolean(); | ||
| FileCompletionListener fileCompletionListener = new FileCompletionListener(numParts, fileName, completionListener); | ||
|
|
||
| for (int partNumber = 0; partNumber < numParts; partNumber++) { | ||
| StreamCompletionListener streamCompletionListener = new StreamCompletionListener( | ||
| partNumber, | ||
| readContext.getPartStreams().get(partNumber), | ||
| fileLocation, | ||
| anyPartStreamFailed, | ||
| fileCompletionListener | ||
| ); | ||
| new ThreadedActionListener<>(logger, threadPool, ThreadPool.Names.GENERIC, streamCompletionListener, false).onResponse(null); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public void onFailure(Exception e) { | ||
| completionListener.onFailure(e); | ||
| } | ||
| } |
94 changes: 94 additions & 0 deletions
94
...n/java/org/opensearch/common/blobstore/stream/read/listener/StreamCompletionListener.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,94 @@ | ||
| /* | ||
| * 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.common.blobstore.stream.read.listener; | ||
|
|
||
| import org.opensearch.common.io.Channels; | ||
| import org.opensearch.common.io.InputStreamContainer; | ||
| import org.opensearch.core.action.ActionListener; | ||
|
|
||
| import java.io.IOException; | ||
| import java.io.InputStream; | ||
| import java.nio.ByteBuffer; | ||
| import java.nio.channels.FileChannel; | ||
| import java.nio.file.Files; | ||
| import java.nio.file.Path; | ||
| import java.nio.file.StandardOpenOption; | ||
| import java.util.concurrent.atomic.AtomicBoolean; | ||
|
|
||
| /** | ||
| * StreamCompletionListener transfers the provided stream into the specified file path using a {@link FileChannel} | ||
| * instance. It performs offset based writes to the file and notifies the {@link FileCompletionListener} on completion. | ||
| * | ||
| * @opensearch.internal | ||
| */ | ||
| public class StreamCompletionListener implements ActionListener<Void> { | ||
| private final int partNumber; | ||
| private final InputStreamContainer blobPartStreamContainer; | ||
| private final Path fileLocation; | ||
| private final AtomicBoolean anyPartStreamFailed; | ||
| private final FileCompletionListener fileCompletionListener; | ||
|
|
||
| // 8 MB buffer for transfer | ||
| private static final int BUFFER_SIZE = 8 * 1024 * 2024; | ||
|
|
||
| public StreamCompletionListener( | ||
| int partNumber, | ||
| InputStreamContainer blobPartStreamContainer, | ||
| Path fileLocation, | ||
| AtomicBoolean anyPartStreamFailed, | ||
| FileCompletionListener fileCompletionListener | ||
| ) { | ||
| this.partNumber = partNumber; | ||
| this.blobPartStreamContainer = blobPartStreamContainer; | ||
| this.fileLocation = fileLocation; | ||
| this.anyPartStreamFailed = anyPartStreamFailed; | ||
| this.fileCompletionListener = fileCompletionListener; | ||
| } | ||
|
|
||
| @Override | ||
| public void onResponse(Void unused) { | ||
| // Ensures no writes to the file if any stream fails. | ||
| if (!anyPartStreamFailed.get()) { | ||
| try (FileChannel outputFileChannel = FileChannel.open(fileLocation, StandardOpenOption.WRITE, StandardOpenOption.CREATE)) { | ||
| try (InputStream inputStream = blobPartStreamContainer.getInputStream()) { | ||
| outputFileChannel.position(blobPartStreamContainer.getOffset()); | ||
|
|
||
| final byte[] buffer = new byte[BUFFER_SIZE]; | ||
| ByteBuffer byteBuffer = ByteBuffer.wrap(buffer); | ||
| int bytesRead; | ||
|
|
||
| while ((bytesRead = inputStream.read(buffer)) != -1) { | ||
| byteBuffer.limit(bytesRead); | ||
| Channels.writeToChannel(byteBuffer, outputFileChannel); | ||
| byteBuffer.clear(); | ||
| } | ||
| } | ||
| } catch (IOException e) { | ||
| onFailure(e); | ||
| return; | ||
| } | ||
| fileCompletionListener.onResponse(partNumber); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public void onFailure(Exception e) { | ||
| try { | ||
| if (Files.exists(fileLocation)) { | ||
| Files.delete(fileLocation); | ||
| } | ||
| } catch (IOException ex) { | ||
| // Die silently | ||
| } | ||
| if (!anyPartStreamFailed.get()) { | ||
| anyPartStreamFailed.compareAndSet(false, true); | ||
| fileCompletionListener.onFailure(e); | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.