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..f57c583a379 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 @@ -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 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..dd025c01fc7 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,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 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..af2ff1798e0 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,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) }