fix: Close threads when calling PutObject (#180)

* fix: Close threads when calling PutObject

---------

Co-authored-by: Lucas McDonald <[email protected]>

Author: lucasmcdonald3

src/main/java/software/amazon/encryption/s3/S3EncryptionClient.java

@@ -192,12 +192,28 @@ public PutObjectResponse putObject(PutObjectRequest putObjectRequest, RequestBod
                 .secureRandom(_secureRandom)
                 .build();
 
+        ExecutorService singleThreadExecutor = Executors.newSingleThreadExecutor();
+
         try {
-            CompletableFuture<PutObjectResponse> futurePut = pipeline.putObject(putObjectRequest, AsyncRequestBody.fromInputStream(requestBody.contentStreamProvider().newStream(), requestBody.optionalContentLength().orElse(-1L), Executors.newSingleThreadExecutor()));
-            return futurePut.join();
+            CompletableFuture<PutObjectResponse> futurePut = pipeline.putObject(putObjectRequest,
+                AsyncRequestBody.fromInputStream(
+                    requestBody.contentStreamProvider().newStream(),
+                    requestBody.optionalContentLength().orElse(-1L),
+                    singleThreadExecutor
+                )
+            );
+
+            PutObjectResponse response = futurePut.join();
+
+            singleThreadExecutor.shutdown();
+
+            return response;
+
         } catch (CompletionException completionException) {
+            singleThreadExecutor.shutdownNow();
             throw new S3EncryptionClientException(completionException.getMessage(), completionException.getCause());
         } catch (Exception exception) {
+            singleThreadExecutor.shutdownNow();
             throw new S3EncryptionClientException(exception.getMessage(), exception);
         }
 
@@ -279,7 +295,6 @@ private CompleteMultipartUploadResponse multipartPutObject(PutObjectRequest requ
         }
 
         ExecutorService es = multipartConfiguration.executorService();
-        final boolean defaultExecutorService = es == null;
         if (es == null) {
             throw new S3EncryptionClientException("ExecutorService should not be null, Please initialize during MultipartConfiguration");
         }
@@ -315,8 +330,8 @@ private CompleteMultipartUploadResponse multipartPutObject(PutObjectRequest requ
         } catch (IOException | InterruptedException | ExecutionException | RuntimeException | Error ex) {
             throw onAbort(observer, ex);
         } finally {
-            if (defaultExecutorService) {
-                // shut down the locally created thread pool
+            if (multipartConfiguration.usingDefaultExecutorService()) {
+                // shut down the thread pool if it was created by the encryption client
                 es.shutdownNow();
             }
             // delete left-over temp files

src/main/java/software/amazon/encryption/s3/internal/MultipartUploadObjectPipeline.java

@@ -32,6 +32,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
 import static software.amazon.encryption.s3.internal.ApiNameVersion.API_NAME_INTERCEPTOR;
@@ -134,10 +135,18 @@ public UploadPartResponse uploadPart(UploadPartRequest request, RequestBody requ
         // Checks the parts are uploaded in series
         materials.beginPartUpload(actualRequest.partNumber(), partContentLength);
         Cipher cipher = materials.getCipher(materials.getIv());
+
+        ExecutorService singleThreadExecutor = Executors.newSingleThreadExecutor();
+
         try {
-            final AsyncRequestBody cipherAsyncRequestBody = new CipherAsyncRequestBody(AsyncRequestBody.fromInputStream(requestBody.contentStreamProvider().newStream(),
+            final AsyncRequestBody cipherAsyncRequestBody = new CipherAsyncRequestBody(
+                AsyncRequestBody.fromInputStream(
+                    requestBody.contentStreamProvider().newStream(),
                     partContentLength, // this MUST be the original contentLength; it refers to the plaintext stream
-                    Executors.newSingleThreadExecutor()), ciphertextLength, materials, cipher.getIV(), isLastPart);
+                    singleThreadExecutor
+                ),
+                ciphertextLength, materials, cipher.getIV(), isLastPart
+            );
 
             // Ensure we haven't already seen the last part
             if (isLastPart) {
@@ -156,6 +165,9 @@ public UploadPartResponse uploadPart(UploadPartRequest request, RequestBody requ
         if (isLastPart) {
             materials.setHasFinalPartBeenSeen(true);
         }
+
+        singleThreadExecutor.shutdown();
+
         return response;
     }
 

src/main/java/software/amazon/encryption/s3/materials/MultipartConfiguration.java

@@ -14,6 +14,7 @@ public class MultipartConfiguration {
     private final long _diskLimit;
     private final UploadObjectObserver _observer;
     private final ExecutorService _es;
+    private final boolean _usingDefaultExecutorService;
     private final MultiFileOutputStream _outputStream;
 
     public MultipartConfiguration(Builder builder) {
@@ -22,6 +23,7 @@ public MultipartConfiguration(Builder builder) {
         this._diskLimit = builder._diskLimit;
         this._observer = builder._observer;
         this._es = builder._es;
+        this._usingDefaultExecutorService = builder._usingDefaultExecutorService;
         this._outputStream = builder._outputStream;
     }
 
@@ -53,6 +55,10 @@ public ExecutorService executorService() {
         return _es;
     }
 
+    public boolean usingDefaultExecutorService() {
+        return _usingDefaultExecutorService;
+    }
+
     static public class Builder {
         private final long MIN_PART_SIZE = 5 << 20;
         private MultiFileOutputStream _outputStream = new MultiFileOutputStream();
@@ -64,6 +70,7 @@ static public class Builder {
         private UploadObjectObserver _observer = new UploadObjectObserver();
         // If null, ExecutorService will be initialized in build() based on maxConnections.
         private ExecutorService _es = null;
+        private boolean _usingDefaultExecutorService;
 
         private Builder() {
         }
@@ -104,6 +111,9 @@ public Builder multiFileOutputStream(MultiFileOutputStream outputStream) {
         public MultipartConfiguration build() {
             if (_es == null) {
                 _es = Executors.newFixedThreadPool(_maxConnections);
+                _usingDefaultExecutorService = true;
+            } else {
+                _usingDefaultExecutorService = false;
             }
 
             return new MultipartConfiguration(this);

src/test/java/software/amazon/encryption/s3/S3EncryptionClientMultipartUploadTest.java

@@ -35,6 +35,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -73,11 +74,14 @@ public void multipartPutObjectAsync() throws IOException {
         Map<String, String> encryptionContext = new HashMap<>();
         encryptionContext.put("user-metadata-key", "user-metadata-value-v3-to-v3");
 
+        ExecutorService singleThreadExecutor = Executors.newSingleThreadExecutor();
+        
         CompletableFuture<PutObjectResponse> futurePut = v3Client.putObject(builder -> builder
                 .bucket(BUCKET)
                 .overrideConfiguration(withAdditionalConfiguration(encryptionContext))
-                .key(objectKey), AsyncRequestBody.fromInputStream(inputStream, fileSizeLimit, Executors.newSingleThreadExecutor()));
+                .key(objectKey), AsyncRequestBody.fromInputStream(inputStream, fileSizeLimit, singleThreadExecutor));
         futurePut.join();
+        singleThreadExecutor.shutdown();
 
         // Asserts
         CompletableFuture<ResponseInputStream<GetObjectResponse>> getFuture = v3Client.getObject(builder -> builder
@@ -110,12 +114,15 @@ public void multipartPutObjectAsyncLargeObjectFails() {
         Map<String, String> encryptionContext = new HashMap<>();
         encryptionContext.put("user-metadata-key", "user-metadata-value-v3-to-v3");
 
+        ExecutorService singleThreadExecutor = Executors.newSingleThreadExecutor();
+
         assertThrows(S3EncryptionClientException.class, () -> v3Client.putObject(builder -> builder
                 .bucket(BUCKET)
                 .overrideConfiguration(withAdditionalConfiguration(encryptionContext))
-                .key(objectKey), AsyncRequestBody.fromInputStream(inputStream, fileSizeLimit, Executors.newSingleThreadExecutor())));
+                .key(objectKey), AsyncRequestBody.fromInputStream(inputStream, fileSizeLimit, singleThreadExecutor)));
 
         v3Client.close();
+        singleThreadExecutor.shutdown();
     }
 
     @Test

src/test/java/software/amazon/encryption/s3/S3EncryptionClientStreamTest.java

@@ -40,6 +40,7 @@
 import java.security.Security;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -324,13 +325,15 @@ public void customSetBufferSizeWithLargeObjectAsyncClient() throws IOException {
         // Tight bound on the custom buffer size limit of 32MiB
         final long fileSizeExceedingDefaultLimit = 1024 * 1024 * 32 + 1;
         final InputStream largeObjectStream = new BoundedInputStream(fileSizeExceedingDefaultLimit);
+        ExecutorService singleThreadExecutor = Executors.newSingleThreadExecutor();
         CompletableFuture<PutObjectResponse> futurePut = v3ClientWithBuffer32MiB.putObject(PutObjectRequest.builder()
                                                                                            .bucket(BUCKET)
                                                                                            .key(objectKey)
-                                                                                           .build(), AsyncRequestBody.fromInputStream(largeObjectStream, fileSizeExceedingDefaultLimit, Executors.newSingleThreadExecutor()));
+                                                                                           .build(), AsyncRequestBody.fromInputStream(largeObjectStream, fileSizeExceedingDefaultLimit, singleThreadExecutor));
 
         futurePut.join();
         largeObjectStream.close();
+        singleThreadExecutor.shutdown();
 
         try {
             // Object is larger than Buffer, so getObject fails