diff --git a/gems/aws-sdk-s3/CHANGELOG.md b/gems/aws-sdk-s3/CHANGELOG.md index 5e22be56f46..9803a5238f6 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.229.0 (2026-08-06) ------------------ 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..dda3f40c1bc 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,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 diff --git a/gems/aws-sdk-s3/lib/aws-sdk-s3/multipart_file_uploader.rb b/gems/aws-sdk-s3/lib/aws-sdk-s3/multipart_file_uploader.rb index 3f29be384c7..96059118201 100644 --- a/gems/aws-sdk-s3/lib/aws-sdk-s3/multipart_file_uploader.rb +++ b/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 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..c68a27c70c1 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 @@ -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 } 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_file_uploader_spec.rb b/gems/aws-sdk-s3/spec/multipart_file_uploader_spec.rb index 23aea0c5800..e168805dd47 100644 --- a/gems/aws-sdk-s3/spec/multipart_file_uploader_spec.rb +++ b/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 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..b6b778040e2 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,86 @@ module S3 end.to raise_error(S3::MultipartUploadError, /failed to abort multipart upload: network-error/) end + it 'aborts without hanging when the executor rejects a task mid-stream' do + client.stub_responses(:create_multipart_upload, upload_id: 'MultipartUploadId') + 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 = MultipartStreamUploader.new(client: client, executor: executor, part_size: 5 * 1024 * 1024) + + expect(client).to receive(:abort_multipart_upload) + .with(params.merge(upload_id: 'MultipartUploadId')).and_call_original + expect do + uploader.upload(params) do |write_stream| + 15.times { write_stream << one_mb } + rescue Errno::EPIPE + # producer stops writing once the read end is closed + end + end.to raise_error(S3::MultipartUploadError) + 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