fix: Address unbounded memory growth in upload_stream when source outpaces upload

gems/aws-sdk-s3/lib/aws-sdk-s3/default_executor.rb

@@ -12,7 +12,7 @@ class DefaultExecutor
       def initialize(options = {})
         @max_threads = options[:max_threads] || DEFAULT_MAX_THREADS
         @state = RUNNING
-        @queue = Queue.new
+        @queue = SizedQueue.new(@max_threads)
         @pool = []
         @mutex = Mutex.new
       end
@@ -25,9 +25,9 @@ def post(*args, &block)
         @mutex.synchronize do
           raise 'Executor has been shutdown and is no longer accepting tasks' unless @state == RUNNING
 
-          @queue << [args, block]
           ensure_worker_available
         end
+        @queue << [args, block]
         true
       end
 

gems/aws-sdk-s3/lib/aws-sdk-s3/multipart_stream_uploader.rb

@@ -113,18 +113,21 @@ def complete_opts(options)
       def read_to_part_body(read_pipe)
         return if read_pipe.closed?
 
-        temp_io = @tempfile ? Tempfile.new('aws-sdk-s3-upload_stream') : StringIO.new(String.new)
-        temp_io.binmode
-        bytes_copied = IO.copy_stream(read_pipe, temp_io, @part_size)
-        temp_io.rewind
-        if bytes_copied.zero?
-          if temp_io.is_a?(Tempfile)
+        if @tempfile
+          temp_io = Tempfile.new('aws-sdk-s3-upload_stream')
+          temp_io.binmode
+          bytes_copied = IO.copy_stream(read_pipe, temp_io, @part_size)
+          temp_io.rewind
+          if bytes_copied.zero?
             temp_io.close
             temp_io.unlink
+            nil
+          else
+            temp_io
           end
-          nil
         else
-          temp_io
+          data = read_pipe.read(@part_size)
+          data.nil? || data.empty? ? nil : StringIO.new(data)
         end
       end
 

gems/aws-sdk-s3/spec/multipart_stream_uploader_spec.rb

@@ -153,6 +153,51 @@ module S3
           end.to raise_error(S3::MultipartUploadError, /failed to abort multipart upload: network-error/)
         end
 
+        context 'when source outpaces upload' do
+          let(:num_threads) { 4 }
+          let(:executor) { DefaultExecutor.new(max_threads: num_threads) }
+          let(:subject) { MultipartStreamUploader.new(client: client, executor: executor, part_size: 1024 * 1024) }
+
+          it 'bounds concurrent in-flight parts to the thread count' do
+            client.stub_responses(:create_multipart_upload, upload_id: 'id')
+            client.stub_responses(:complete_multipart_upload)
+            mutex = Mutex.new
+            in_flight = 0
+            peak_in_flight = 0
+            allow(client).to receive(:upload_part) do |_part|
+              mutex.synchronize do
+                in_flight += 1
+                peak_in_flight = in_flight if in_flight > peak_in_flight
+              end
+              sleep(0.05)
+              mutex.synchronize { in_flight -= 1 }
+            end.and_return(double(:upload_part, etag: 'etag'))
+
+            subject.upload(params) do |write_stream|
+              20.times { write_stream << one_mb }
+            end
+
+            expect(peak_in_flight).to be <= num_threads
+          end
+
+          it 'completes all parts under backpressure' do
+            client.stub_responses(:create_multipart_upload, upload_id: 'id')
+            client.stub_responses(:complete_multipart_upload)
+            mutex = Mutex.new
+            uploaded_parts = []
+            allow(client).to receive(:upload_part) do |part|
+              sleep(0.05)
+              mutex.synchronize { uploaded_parts << part[:part_number] }
+            end.and_return(double(:upload_part, etag: 'etag'))
+
+            subject.upload(params) do |write_stream|
+              10.times { write_stream << one_mb }
+            end
+
+            expect(uploaded_parts.sort).to eq((1..10).to_a)
+          end
+        end
+
         context 'when tempfile is true' do
           let(:subject) { MultipartStreamUploader.new(client: client, tempfile: true, executor: DefaultExecutor.new) }