Merge pull request DSpace#9477 from 4Science/lazy_s3input

Improve S3 Bitstream Storage to Lazy download object from S3

dspace-api/src/main/java/org/dspace/storage/bitstore/S3BitStoreService.java

@@ -10,6 +10,7 @@
 import static java.lang.String.valueOf;
 
 import java.io.File;
+import java.io.FileNotFoundException;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -102,6 +103,11 @@ public class S3BitStoreService extends BaseBitStoreService {
  private String awsRegionName;
  private boolean useRelativePath;
 
+ /**
+ * The maximum size of individual chunk to download from S3 when a file is accessed. Default 5Mb
+ */
+ private long bufferSize = 5 * 1024 * 1024;
+
  /**
  * container for all the assets
  */
@@ -258,20 +264,7 @@ public InputStream get(Bitstream bitstream) throws IOException {
  if (isRegisteredBitstream(key)) {
  key = key.substring(REGISTERED_FLAG.length());
  }
- try {
- File tempFile = File.createTempFile("s3-disk-copy-" + UUID.randomUUID(), "temp");
- tempFile.deleteOnExit();
-
- GetObjectRequest getObjectRequest = new GetObjectRequest(bucketName, key);
-
- Download download = tm.download(getObjectRequest, tempFile);
- download.waitForCompletion();
-
- return new DeleteOnCloseFileInputStream(tempFile);
- } catch (AmazonClientException | InterruptedException e) {
- log.error("get(" + key + ")", e);
- throw new IOException(e);
- }
+ return new S3LazyInputStream(key, bufferSize, bitstream.getSizeBytes());
  }
 
  /**
@@ -622,4 +615,84 @@ public boolean isRegisteredBitstream(String internalId) {
  return internalId.startsWith(REGISTERED_FLAG);
  }
 
+ public void setBufferSize(long bufferSize) {
+ this.bufferSize = bufferSize;
+ }
+
+ /**
+ * This inner class represent an InputStream that uses temporary files to
+ * represent chunk of the object downloaded from S3. When the input stream is
+ * read the class look first to the current chunk and download a new one once if
+ * the current one as been fully read. The class is responsible to close a chunk
+ * as soon as a new one is retrieved, the last chunk is closed when the input
+ * stream itself is closed or the last byte is read (the first of the two)
+ */
+ public class S3LazyInputStream extends InputStream {
+ private InputStream currentChunkStream;
+ private String objectKey;
+ private long endOfChunk = -1;
+ private long chunkMaxSize;
+ private long currPos = 0;
+ private long fileSize;
+
+ public S3LazyInputStream(String objectKey, long chunkMaxSize, long fileSize) throws IOException {
+ this.objectKey = objectKey;
+ this.chunkMaxSize = chunkMaxSize;
+ this.endOfChunk = 0;
+ this.fileSize = fileSize;
+ downloadChunk();
+ }
+
+ @Override
+ public int read() throws IOException {
+ // is the current chunk completely read and other are available?
+ if (currPos == endOfChunk && currPos < fileSize) {
+ currentChunkStream.close();
+ downloadChunk();
+ }
+
+ int byteRead = currPos < endOfChunk ? currentChunkStream.read() : -1;
+ // do we get any data or are we at the end of the file?
+ if (byteRead != -1) {
+ currPos++;
+ } else {
+ currentChunkStream.close();
+ }
+ return byteRead;
+ }
+
+ /**
+ * This method download the next chunk from S3
+ *
+ * @throws IOException
+ * @throws FileNotFoundException
+ */
+ private void downloadChunk() throws IOException, FileNotFoundException {
+ // Create a DownloadFileRequest with the desired byte range
+ long startByte = currPos; // Start byte (inclusive)
+ long endByte = Long.min(startByte + chunkMaxSize - 1, fileSize - 1); // End byte (inclusive)
+ GetObjectRequest getRequest = new GetObjectRequest(bucketName, objectKey)
+ .withRange(startByte, endByte);
+
+ File currentChunkFile = File.createTempFile("s3-disk-copy-" + UUID.randomUUID(), "temp");
+ currentChunkFile.deleteOnExit();
+ try {
+ Download download = tm.download(getRequest, currentChunkFile);
+ download.waitForCompletion();
+ currentChunkStream = new DeleteOnCloseFileInputStream(currentChunkFile);
+ endOfChunk = endOfChunk + download.getProgress().getBytesTransferred();
+ } catch (AmazonClientException | InterruptedException e) {
+ currentChunkFile.delete();
+ throw new IOException(e);
+ }
+ }
+
+ @Override
+ public void close() throws IOException {
+ if (currentChunkStream != null) {
+ currentChunkStream.close();
+ }
+ }
+
+ }
 }

dspace-api/src/test/java/org/dspace/storage/bitstore/S3BitStoreServiceIT.java

@@ -99,6 +99,7 @@ public void setup() throws Exception {
  s3BitStoreService = new S3BitStoreService(amazonS3Client);
  s3BitStoreService.setEnabled(BooleanUtils.toBoolean(
  configurationService.getProperty("assetstore.s3.enabled")));
+ s3BitStoreService.setBufferSize(22);
  context.turnOffAuthorisationSystem();
 
  parentCommunity = CommunityBuilder.createCommunity(context)
@@ -129,12 +130,25 @@ public void testBitstreamPutAndGetWithAlreadyPresentBucket() throws IOException
  assertThat(amazonS3Client.listBuckets(), contains(bucketNamed(bucketName)));
 
  context.turnOffAuthorisationSystem();
- String content = "Test bitstream content";
+ String content = "Test bitstream content";
+ String contentOverOneSpan = "This content span two chunks";
+ String contentExactlyTwoSpans = "Test bitstream contentTest bitstream content";
+ String contentOverOneTwoSpans = "Test bitstream contentThis content span three chunks";
  Bitstream bitstream = createBitstream(content);
+ Bitstream bitstreamOverOneSpan = createBitstream(contentOverOneSpan);
+ Bitstream bitstreamExactlyTwoSpans = createBitstream(contentExactlyTwoSpans);
+ Bitstream bitstreamOverOneTwoSpans = createBitstream(contentOverOneTwoSpans);
  context.restoreAuthSystemState();
 
- s3BitStoreService.put(bitstream, toInputStream(content));
+ checkGetPut(bucketName, content, bitstream);
+ checkGetPut(bucketName, contentOverOneSpan, bitstreamOverOneSpan);
+ checkGetPut(bucketName, contentExactlyTwoSpans, bitstreamExactlyTwoSpans);
+ checkGetPut(bucketName, contentOverOneTwoSpans, bitstreamOverOneTwoSpans);
+
+ }
 
+ private void checkGetPut(String bucketName, String content, Bitstream bitstream) throws IOException {
+ s3BitStoreService.put(bitstream, toInputStream(content));
  String expectedChecksum = Utils.toHex(generateChecksum(content));
 
  assertThat(bitstream.getSizeBytes(), is((long) content.length()));
@@ -147,7 +161,6 @@ public void testBitstreamPutAndGetWithAlreadyPresentBucket() throws IOException
  String key = s3BitStoreService.getFullKey(bitstream.getInternalId());
  ObjectMetadata objectMetadata = amazonS3Client.getObjectMetadata(bucketName, key);
  assertThat(objectMetadata.getContentMD5(), is(expectedChecksum));
-
  }
 
  @Test