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

gems/aws-sdk-s3/CHANGELOG.md

@@ -1,6 +1,8 @@
 Unreleased Changes
 ------------------
 
+* Issue - Bound memory usage in `upload_stream` when the source produces data faster than parts can be uploaded.
+
 1.229.0 (2026-08-06)
 ------------------
 

gems/aws-sdk-s3/lib/aws-sdk-s3/customizations/object.rb

@@ -380,7 +380,9 @@ def public_url(options = {})
       #   and {Client#upload_part} can be provided.
       #
       # @option options [Integer] :thread_count (10) The number of parallel multipart uploads.
-      #   An additional thread is used internally for task coordination.
+      #   An additional thread is used internally for task coordination. This also bounds
+      #   how many parts are buffered ahead of the upload, limiting memory usage to roughly
+      #   `2 * :thread_count * :part_size`.
       #
       # @option options [Boolean] :tempfile (false) Normally read data is stored
       #   in memory when building the parts in order to complete the underlying
@@ -405,7 +407,10 @@ def public_url(options = {})
       # @see Client#upload_part
       def upload_stream(options = {}, &block)
         upload_opts = options.merge(bucket: bucket_name, key: key)
-        executor = DefaultExecutor.new(max_threads: upload_opts.delete(:thread_count))
+        thread_count = upload_opts.delete(:thread_count) || DefaultExecutor::DEFAULT_MAX_THREADS
+        # A bounded queue prevents the source from reading ahead without limit when it
+        # produces data faster than parts can be uploaded.
+        executor = DefaultExecutor.new(max_threads: thread_count, max_queue: thread_count)
         uploader = MultipartStreamUploader.new(
           client: client,
           executor: executor,

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

@@ -11,8 +11,9 @@ class DefaultExecutor
 
       def initialize(options = {})
         @max_threads = options[:max_threads] || DEFAULT_MAX_THREADS
+        @max_queue = options[:max_queue] || 0
         @state = RUNNING
-        @queue = Queue.new
+        @queue = @max_queue.zero? ? Queue.new : SizedQueue.new(@max_queue) # 0 is unbounded
         @pool = []
         @mutex = Mutex.new
       end
@@ -25,10 +26,15 @@ 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
+        # Pushed outside the mutex because a bounded queue blocks the caller when
+        # full and holding the lock while parked would deadlock #shutdown and #kill.
+        @queue.push([args, block])
         true
+      rescue ClosedQueueError
+        # shutdown or kill happened while parked on a full queue
+        raise 'Executor has been shutdown and is no longer accepting tasks'
       end
 
       # Immediately terminates all worker threads and clears pending tasks.
@@ -38,6 +44,7 @@ def post(*args, &block)
       def kill
         @mutex.synchronize do
           @state = SHUTDOWN
+          @queue.close # wakes any producer parked on a full queue
           @pool.each(&:kill)
           @pool.clear
           @queue.clear
@@ -56,7 +63,9 @@ def shutdown(timeout = nil)
           return true if @state == SHUTDOWN
 
           @state = SHUTTING_DOWN
-          @pool.size.times { @queue << :shutdown }
+          # Closing wakes parked producers and lets workers drain remaining tasks
+          # before exiting without pushing sentinels onto a queue that may be full.
+          @queue.close
         end
 
         if timeout
@@ -91,8 +100,6 @@ def ensure_worker_available
       def spawn_worker
         Thread.new do
           while (job = @queue.shift)
-            break if job == :shutdown
-
             args, block = job
             block.call(*args)
           end

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

@@ -149,21 +149,32 @@ def upload_with_executor(pending, completed, options)
         while (part = pending.shift)
           break if abort_upload
 
-          upload_attempts += 1
-          @executor.post(part) do |p|
-            Thread.current[:net_http_override_body_stream_chunk] = @http_chunk_size if @http_chunk_size
-            update_progress(progress, p)
-            resp = @client.upload_part(p)
-            completed_part = { etag: resp.etag, part_number: p[:part_number] }
-            apply_part_checksum(resp, completed_part)
-            completed.push(completed_part)
+          begin
+            @executor.post(part) do |p|
+              Thread.current[:net_http_override_body_stream_chunk] = @http_chunk_size if @http_chunk_size
+              update_progress(progress, p)
+              resp = @client.upload_part(p)
+              completed_part = { etag: resp.etag, part_number: p[:part_number] }
+              apply_part_checksum(resp, completed_part)
+              completed.push(completed_part)
+            rescue StandardError => e
+              abort_upload = true
+              errors << e
+            ensure
+              p[:body].close
+              Thread.current[:net_http_override_body_stream_chunk] = nil if @http_chunk_size
+              completion_queue << :done
+            end
+            # Count only successfully queued parts; a failed post never runs the
+            # block, so it never pushes :done and must not be waited on below.
+            upload_attempts += 1
           rescue StandardError => e
+            # The executor rejected the task (e.g. shut down mid-upload). Record
+            # it so the abort ceremony runs instead of the error escaping and
+            # leaving the multipart upload orphaned on S3.
             abort_upload = true
             errors << e
-          ensure
-            p[:body].close
-            Thread.current[:net_http_override_body_stream_chunk] = nil if @http_chunk_size
-            completion_queue << :done
+            break
           end
         end
 

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

@@ -113,18 +113,24 @@ 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
+          # Read into a single right-sized buffer. IO.copy_stream into a StringIO grows
+          # the backing string geometrically (an 8MB buffer for a 5MB part) and discards
+          # the intermediates, fragmenting the heap across concurrent parts.
+          data = read_pipe.read(@part_size)
+          data.nil? ? nil : StringIO.new(data)
         end
       end
 
@@ -139,20 +145,34 @@ def upload_with_executor(read_pipe, completed, errors, options)
           end
           break unless part_body || current_part_num == 1
 
-          queued_parts += 1
-          @executor.post(part_body, current_part_num, options) do |body, num, opts|
-            part = opts.merge(body: body, part_number: num)
-            resp = @client.upload_part(part)
-            completed_part = create_completed_part(resp, part)
-            completed.push(completed_part)
+          begin
+            @executor.post(part_body, current_part_num, options) do |body, num, opts|
+              part = opts.merge(body: body, part_number: num)
+              resp = @client.upload_part(part)
+              completed_part = create_completed_part(resp, part)
+              completed.push(completed_part)
+            rescue StandardError => e
+              mutex.synchronize do
+                errors.push(e)
+                read_pipe.close_read unless read_pipe.closed?
+              end
+            ensure
+              clear_body(body)
+              completion_queue << :done
+            end
+            # Count only successfully queued parts; a failed post never runs the
+            # block, so it never pushes :done and must not be waited on below.
+            queued_parts += 1
           rescue StandardError => e
+            # The executor rejected the task (e.g. shut down mid-stream). Record
+            # the error and close the read end so the producer block stops writing
+            # instead of blocking forever on a full pipe, letting the abort run.
             mutex.synchronize do
               errors.push(e)
               read_pipe.close_read unless read_pipe.closed?
             end
-          ensure
-            clear_body(body)
-            completion_queue << :done
+            clear_body(part_body)
+            break
           end
         end
         queued_parts.times { completion_queue.pop }

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

@@ -491,6 +491,9 @@ def upload_file(source, bucket:, key:, **options)
       # @option options [Integer] :thread_count (10)
       #   The number of parallel multipart uploads. Only used when no custom executor is provided (creates
       #   {DefaultExecutor} with the given thread count). An additional thread is used internally for task coordination.
+      #   This also bounds how many parts are buffered ahead of the upload, limiting memory usage to roughly
+      #   `2 * :thread_count * :part_size`. When a custom `:executor` is provided, it is responsible for applying
+      #   its own backpressure.
       #
       # @option options [Boolean] :tempfile (false)
       #   Normally read data is stored in memory when building the parts in order to complete the underlying
@@ -511,7 +514,10 @@ def upload_file(source, bucket:, key:, **options)
       # @see Client#upload_part
       def upload_stream(bucket:, key:, **options, &block)
         upload_opts = options.merge(bucket: bucket, key: key)
-        executor = @executor || DefaultExecutor.new(max_threads: upload_opts.delete(:thread_count))
+        thread_count = upload_opts.delete(:thread_count) || DefaultExecutor::DEFAULT_MAX_THREADS
+        # A bounded queue prevents the source from reading ahead without limit when it
+        # produces data faster than parts can be uploaded.
+        executor = @executor || DefaultExecutor.new(max_threads: thread_count, max_queue: thread_count)
         uploader = MultipartStreamUploader.new(
           client: @client,
           executor: executor,

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

@@ -24,6 +24,62 @@ module S3
         end
       end
 
+      context 'when the queue is full' do
+        let(:executor) { DefaultExecutor.new(max_threads: 1, max_queue: 1) }
+        let(:release) { Queue.new }
+        let(:errors) { [] }
+
+        # occupies the only worker, then fills the single queue slot
+        def fill_queue
+          started = Queue.new
+          executor.post do
+            started << :running
+            release.pop
+          end
+          started.pop
+          executor.post {}
+        end
+
+        def park_producer
+          fill_queue
+          producer = Thread.new do
+            executor.post {}
+          rescue RuntimeError => e
+            errors << e
+          end
+          sleep(0.1)
+          raise 'producer did not park' unless producer.status == 'sleep'
+
+          producer
+        end
+
+        it 'blocks the caller until a worker frees a slot' do
+          fill_queue
+          parked = Thread.new { executor.post {} }
+          sleep(0.1)
+          expect(parked.status).to eq('sleep')
+
+          release << :go
+          expect(parked.value).to be(true)
+          executor.shutdown
+        end
+
+        it 'kill unblocks the producer instead of silently dropping the task' do
+          producer = park_producer
+          expect(executor.kill).to be(true)
+          expect(producer.join(2)).to_not be_nil
+          expect(errors.first).to be_a(RuntimeError)
+        end
+
+        it 'shutdown does not deadlock while holding the lock' do
+          producer = park_producer
+          shutdown = Thread.new { executor.shutdown(0.5) }
+          expect(shutdown.join(2)).to_not be_nil
+          expect(producer.join(2)).to_not be_nil
+          expect(errors.first).to be_a(RuntimeError)
+        end
+      end
+
       describe '#shutdown' do
         it 'waits for running tasks to be complete' do
           result = nil

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

@@ -170,6 +170,25 @@ module S3
           expect(client).to receive(:abort_multipart_upload).with(params.merge(upload_id: 'MultipartUploadId'))
           expect { subject.upload(large_file, params) }.to raise_error(Aws::S3::MultipartUploadError)
         end
+
+        it 'aborts multipart upload when the executor rejects a task mid-upload' do
+          client.stub_responses(:upload_part, etag: 'etag')
+          executor = DefaultExecutor.new
+          calls = 0
+          # Simulate a concurrent shutdown closing the queue: the second post is
+          # rejected the way DefaultExecutor#post now raises on a closed queue.
+          allow(executor).to receive(:post).and_wrap_original do |original, *args, &blk|
+            calls += 1
+            raise 'Executor has been shutdown and is no longer accepting tasks' if calls == 2
+
+            original.call(*args, &blk)
+          end
+          uploader = MultipartFileUploader.new(client: client, executor: executor)
+
+          expect(client).to receive(:abort_multipart_upload)
+            .with(params.merge(upload_id: 'MultipartUploadId')).and_call_original
+          expect { uploader.upload(large_file, params) }.to raise_error(Aws::S3::MultipartUploadError)
+        end
       end
     end
   end