Merge pull request #1596 from CMSgov/QPPA-11251-memory-leak-mitigation

QPPA-11251: Adds closure to s3 connections to mitigate memory leak issue

Author: john-manack

rest-api/src/main/java/gov/cms/qpp/conversion/api/services/internal/StorageServiceImpl.java

@@ -20,6 +20,8 @@
 import gov.cms.qpp.conversion.api.model.Constants;
 import gov.cms.qpp.conversion.api.services.StorageService;
 
+import java.io.FilterInputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
@@ -99,12 +101,11 @@ public InputStream getFileByLocationId(String fileLocationId) {
 		API_LOG.info("Retrieving file {} from bucket {}", fileLocationId, bucketName);
 
 		GetObjectRequest getObjectRequest = new GetObjectRequest(bucketName, fileLocationId);
-
-		S3Object s3Object = amazonS3.getObject(getObjectRequest);
+		InputStream objectContent = retrieveManagedContentStream(getObjectRequest);
 
 		API_LOG.info("Successfully retrieved file {} from S3 bucket {}", getObjectRequest.getKey(), getObjectRequest.getBucketName());
 
-		return s3Object.getObjectContent();
+		return objectContent;
 	}
 
 	/**
@@ -123,9 +124,7 @@ public InputStream getCpcPlusValidationFile() {
 		API_LOG.info("Retrieving CPC+ validation file from bucket {}", bucketName);
 
 		GetObjectRequest getObjectRequest = new GetObjectRequest(bucketName, key);
-		S3Object s3Object = amazonS3.getObject(getObjectRequest);
-
-		return s3Object.getObjectContent();
+		return retrieveManagedContentStream(getObjectRequest);
 	}
 
 	/**
@@ -143,9 +142,15 @@ public InputStream getApmValidationFile(String fileName) {
 		API_LOG.info("Retrieving APM validation file from bucket {}", bucketName);
 
 		GetObjectRequest getObjectRequest = new GetObjectRequest(bucketName, fileName);
-		S3Object s3Object = amazonS3.getObject(getObjectRequest);
+		return retrieveManagedContentStream(getObjectRequest);
+	}
 
-		return s3Object.getObjectContent();
+	/**
+	 * Provides an {@link InputStream} whose close also closes the backing {@link S3Object}.
+	 */
+	private InputStream retrieveManagedContentStream(GetObjectRequest getObjectRequest) {
+		S3Object s3Object = amazonS3.getObject(getObjectRequest);
+		return new ManagedS3InputStream(s3Object);
 	}
 
 	/**
@@ -176,4 +181,22 @@ protected String asynchronousAction(Supplier<PutObjectRequest> objectToActOn) {
 	protected String getActionName() {
 		return "Write to Storage";
 	}
+
+	private static final class ManagedS3InputStream extends FilterInputStream {
+		private final S3Object s3Object;
+
+		ManagedS3InputStream(S3Object s3Object) {
+			super(s3Object.getObjectContent());
+			this.s3Object = s3Object;
+		}
+
+		@Override
+		public void close() throws IOException {
+			try {
+				super.close();
+			} finally {
+				s3Object.close();
+			}
+		}
+	}
 }

rest-api/src/test/java/gov/cms/qpp/conversion/api/services/internal/StorageServiceImplTest.java

@@ -12,7 +12,10 @@
 import static org.mockito.Mockito.when;
 
 import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 
@@ -144,7 +147,8 @@ void noBucket() {
 	@Test
 	void envVariablesPresent() {
 		S3Object s3ObjectMock = mock(S3Object.class);
-		s3ObjectMock.setObjectContent(new ByteArrayInputStream("1234".getBytes()));
+		S3ObjectInputStream objectContent = new S3ObjectInputStream(new ByteArrayInputStream("1234".getBytes(StandardCharsets.UTF_8)), null);
+		when(s3ObjectMock.getObjectContent()).thenReturn(objectContent);
 		Mockito.when(amazonS3Client.getObject(any(GetObjectRequest.class))).thenReturn(s3ObjectMock);
 		Mockito.when(environment.getProperty(Constants.BUCKET_NAME_ENV_VARIABLE)).thenReturn("meep");
 		underTest.getFileByLocationId("meep");
@@ -197,31 +201,53 @@ void test_getCpcPlusValidationFile_NPE() {
 	}
 
 	@Test
-	void test_getCpcPlusValidationFile() {
-		S3ObjectInputStream expected = new S3ObjectInputStream(null, null);
+	void test_getCpcPlusValidationFile() throws IOException {
+		byte[] expectedBytes = "Mock Contents".getBytes(StandardCharsets.UTF_8);
+		S3ObjectInputStream expectedStream = new S3ObjectInputStream(new ByteArrayInputStream(expectedBytes), null);
 		S3Object mockS3Obj = mock(S3Object.class);
-		Mockito.when(mockS3Obj.getObjectContent()).thenReturn(expected);
+		Mockito.when(mockS3Obj.getObjectContent()).thenReturn(expectedStream);
 
 		Mockito.when(environment.getProperty(Constants.CPC_PLUS_BUCKET_NAME_VARIABLE)).thenReturn("Mock_Bucket");
 		Mockito.when(environment.getProperty(Constants.CPC_PLUS_FILENAME_VARIABLE)).thenReturn("Mock_Filename");
-		Mockito.when(amazonS3Client.getObject( any(GetObjectRequest.class) )).thenReturn(mockS3Obj);
+		Mockito.when(amazonS3Client.getObject(any(GetObjectRequest.class))).thenReturn(mockS3Obj);
 
-		InputStream actual = underTest.getCpcPlusValidationFile();
+		byte[] actualBytes;
+		try (InputStream actual = underTest.getCpcPlusValidationFile()) {
+			assertThat(actual).isNotNull();
+			actualBytes = toByteArray(actual);
+		}
 
-		assertThat(actual).isEqualTo(expected);
+		assertThat(actualBytes).isEqualTo(expectedBytes);
+		verify(mockS3Obj, times(1)).close();
 	}
 
 	@Test
-	void test_getApmValidationFile() {
-		S3ObjectInputStream expected = new S3ObjectInputStream(null, null);
+	void test_getApmValidationFile() throws IOException {
+		byte[] expectedBytes = "APM".getBytes(StandardCharsets.UTF_8);
+		S3ObjectInputStream expectedStream = new S3ObjectInputStream(new ByteArrayInputStream(expectedBytes), null);
 		S3Object mockS3Obj = mock(S3Object.class);
-		Mockito.when(mockS3Obj.getObjectContent()).thenReturn(expected);
+		Mockito.when(mockS3Obj.getObjectContent()).thenReturn(expectedStream);
 
 		Mockito.when(environment.getProperty(Constants.BUCKET_NAME_ENV_VARIABLE)).thenReturn("Mock_Bucket");
-		Mockito.when(amazonS3Client.getObject( any(GetObjectRequest.class) )).thenReturn(mockS3Obj);
+		Mockito.when(amazonS3Client.getObject(any(GetObjectRequest.class))).thenReturn(mockS3Obj);
 
-		InputStream actual = underTest.getApmValidationFile(Constants.CPC_PLUS_APM_FILE_NAME_KEY);
+		byte[] actualBytes;
+		try (InputStream actual = underTest.getApmValidationFile(Constants.CPC_PLUS_APM_FILE_NAME_KEY)) {
+			assertThat(actual).isNotNull();
+			actualBytes = toByteArray(actual);
+		}
 
-		assertThat(actual).isEqualTo(expected);
+		assertThat(actualBytes).isEqualTo(expectedBytes);
+		verify(mockS3Obj, times(1)).close();
+	}
+
+	private byte[] toByteArray(InputStream inputStream) throws IOException {
+		ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+		byte[] data = new byte[1024];
+		int bytesRead;
+		while ((bytesRead = inputStream.read(data)) != -1) {
+			buffer.write(data, 0, bytesRead);
+		}
+		return buffer.toByteArray();
 	}
 }