diff --git a/gems/aws-sdk-s3/CHANGELOG.md b/gems/aws-sdk-s3/CHANGELOG.md index 2385d906d8e..f5609e7e91e 100644 --- a/gems/aws-sdk-s3/CHANGELOG.md +++ b/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.228.1 (2026-07-23) ------------------ diff --git a/gems/aws-sdk-s3/lib/aws-sdk-s3/customizations/object.rb b/gems/aws-sdk-s3/lib/aws-sdk-s3/customizations/object.rb index dd2cb036edf..99b66013c4a 100644 --- a/gems/aws-sdk-s3/lib/aws-sdk-s3/customizations/object.rb +++ b/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, diff --git a/gems/aws-sdk-s3/lib/aws-sdk-s3/default_executor.rb b/gems/aws-sdk-s3/lib/aws-sdk-s3/default_executor.rb index 13a719f4397..8c16f393ed3 100644 --- a/gems/aws-sdk-s3/lib/aws-sdk-s3/default_executor.rb +++ b/gems/aws-sdk-s3/lib/aws-sdk-s3/default_executor.rb @@ -11,8 +11,10 @@ class DefaultExecutor def initialize(options = {}) @max_threads = options[:max_threads] || DEFAULT_MAX_THREADS + @max_queue = options[:max_queue] || 0 @state = RUNNING - @queue = Queue.new + # A bounded queue applies backpressure to producers when full. 0 means unbounded. + @queue = @max_queue.zero? ? Queue.new : SizedQueue.new(@max_queue) @pool = [] @mutex = Mutex.new end @@ -25,10 +27,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 +45,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 +64,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 +101,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 diff --git a/gems/aws-sdk-s3/lib/aws-sdk-s3/multipart_stream_uploader.rb b/gems/aws-sdk-s3/lib/aws-sdk-s3/multipart_stream_uploader.rb index ae6f75a47a1..812dbacbac7 100644 --- a/gems/aws-sdk-s3/lib/aws-sdk-s3/multipart_stream_uploader.rb +++ b/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 diff --git a/gems/aws-sdk-s3/lib/aws-sdk-s3/transfer_manager.rb b/gems/aws-sdk-s3/lib/aws-sdk-s3/transfer_manager.rb index 864dd543c14..a3a49597e25 100644 --- a/gems/aws-sdk-s3/lib/aws-sdk-s3/transfer_manager.rb +++ b/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, diff --git a/gems/aws-sdk-s3/spec/default_executor_spec.rb b/gems/aws-sdk-s3/spec/default_executor_spec.rb index ffb6974d096..27e2aa889e6 100644 --- a/gems/aws-sdk-s3/spec/default_executor_spec.rb +++ b/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 diff --git a/gems/aws-sdk-s3/spec/multipart_stream_uploader_spec.rb b/gems/aws-sdk-s3/spec/multipart_stream_uploader_spec.rb index 6262a42b560..28d75fbbb3e 100644 --- a/gems/aws-sdk-s3/spec/multipart_stream_uploader_spec.rb +++ b/gems/aws-sdk-s3/spec/multipart_stream_uploader_spec.rb @@ -153,6 +153,60 @@ 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) { 2 } + let(:executor) { DefaultExecutor.new(max_threads: num_threads, max_queue: num_threads) } + let(:subject) { MultipartStreamUploader.new(client: client, executor: executor, part_size: 1024 * 1024) } + + it 'bounds the number of parts buffered ahead of the upload' do + client.stub_responses(:create_multipart_upload, upload_id: 'id') + client.stub_responses(:complete_multipart_upload) + mutex = Mutex.new + buffered = 0 + peak_buffered = 0 + # count parts read off the pipe but not yet uploaded + allow(subject).to receive(:read_to_part_body).and_wrap_original do |original, *args| + body = original.call(*args) + mutex.synchronize do + if body + buffered += 1 + peak_buffered = buffered if buffered > peak_buffered + end + end + body + end + allow(client).to receive(:upload_part) do |_part| + sleep(0.05) + mutex.synchronize { buffered -= 1 } + end.and_return(double(:upload_part, etag: 'etag')) + + subject.upload(params) do |write_stream| + 30.times { write_stream << one_mb } + end + + # at most max_queue queued + max_threads in flight + 1 being read. + # without a bounded queue all 30 parts are read into memory up front. + expect(peak_buffered).to be <= (num_threads * 2) + 1 + 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) } diff --git a/gems/aws-sdk-s3/spec/object/upload_stream_spec.rb b/gems/aws-sdk-s3/spec/object/upload_stream_spec.rb index 8f2a5739813..f8683144930 100644 --- a/gems/aws-sdk-s3/spec/object/upload_stream_spec.rb +++ b/gems/aws-sdk-s3/spec/object/upload_stream_spec.rb @@ -29,7 +29,10 @@ module S3 custom_thread_count = 20 client.stub_responses(:create_multipart_upload, upload_id: 'id') client.stub_responses(:complete_multipart_upload) - expect(DefaultExecutor).to receive(:new).with(max_threads: custom_thread_count).and_call_original + expect(DefaultExecutor) + .to receive(:new) + .with(max_threads: custom_thread_count, max_queue: custom_thread_count) + .and_call_original subject.upload_stream(thread_count: custom_thread_count) { |_write_stream| } end