aws-sdk-s3 1.10.0 → 1.208.0

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

@@ -1,49 +1,46 @@
+# frozen_string_literal: true
+
 require 'pathname'
-require 'thread'
 require 'set'
 
 module Aws
   module S3
     # @api private
     class MultipartFileUploader
-
       MIN_PART_SIZE = 5 * 1024 * 1024 # 5MB
-
-      FILE_TOO_SMALL = "unable to multipart upload files smaller than 5MB"
-
       MAX_PARTS = 10_000
-
-      THREAD_COUNT = 10
-
-      # @api private
-      CREATE_OPTIONS =
-        Set.new(Client.api.operation(:create_multipart_upload).input.shape.member_names)
-
-      # @api private
-      UPLOAD_PART_OPTIONS =
-        Set.new(Client.api.operation(:upload_part).input.shape.member_names)
+      CREATE_OPTIONS = Set.new(Client.api.operation(:create_multipart_upload).input.shape.member_names)
+      COMPLETE_OPTIONS = Set.new(Client.api.operation(:complete_multipart_upload).input.shape.member_names)
+      UPLOAD_PART_OPTIONS = Set.new(Client.api.operation(:upload_part).input.shape.member_names)
+      CHECKSUM_KEYS = Set.new(
+        Client.api.operation(:upload_part).input.shape.members.map do |n, s|
+          n if s.location == 'header' && s.location_name.start_with?('x-amz-checksum-')
+        end.compact
+      )
 
       # @option options [Client] :client
       def initialize(options = {})
         @client = options[:client] || Client.new
-        @thread_count = options[:thread_count] || THREAD_COUNT
+        @executor = options[:executor]
       end
 
       # @return [Client]
       attr_reader :client
 
-      # @param [String,Pathname,File,Tempfile] source
-      # @option options [required,String] :bucket
-      # @option options [required,String] :key
-      # @return [void]
+      # @param [String, Pathname, File, Tempfile] source The file to upload.
+      # @option options [required, String] :bucket The bucket to upload to.
+      # @option options [required, String] :key The key for the object.
+      # @option options [Proc] :progress_callback
+      #   A Proc that will be called when each chunk of the upload is sent.
+      #   It will be invoked with [bytes_read], [total_sizes]
+      # @return [Seahorse::Client::Response] - the CompleteMultipartUploadResponse
       def upload(source, options = {})
-        if File.size(source) < MIN_PART_SIZE
-          raise ArgumentError, FILE_TOO_SMALL
-        else
-          upload_id = initiate_upload(options)
-          parts = upload_parts(upload_id, source, options)
-          complete_upload(upload_id, parts, options)
-        end
+        file_size = File.size(source)
+        raise ArgumentError, 'unable to multipart upload files smaller than 5MB' if file_size < MIN_PART_SIZE
+
+        upload_id = initiate_upload(options)
+        parts = upload_parts(upload_id, source, file_size, options)
+        complete_upload(upload_id, parts, file_size, options)
       end
 
       private
@@ -52,18 +49,21 @@ module Aws
         @client.create_multipart_upload(create_opts(options)).upload_id
       end
 
-      def complete_upload(upload_id, parts, options)
+      def complete_upload(upload_id, parts, file_size, options)
         @client.complete_multipart_upload(
-          bucket: options[:bucket],
-          key: options[:key],
+          **complete_opts(options),
           upload_id: upload_id,
-          multipart_upload: { parts: parts })
+          multipart_upload: { parts: parts },
+          mpu_object_size: file_size
+        )
+      rescue StandardError => e
+        abort_upload(upload_id, options, [e])
       end
 
-      def upload_parts(upload_id, source, options)
-        pending = PartList.new(compute_parts(upload_id, source, options))
+      def upload_parts(upload_id, source, file_size, options)
         completed = PartList.new
-        errors = upload_in_threads(pending, completed)
+        pending = PartList.new(compute_parts(upload_id, source, file_size, options))
+        errors = upload_with_executor(pending, completed, options)
         if errors.empty?
           completed.to_a.sort_by { |part| part[:part_number] }
         else
@@ -72,81 +72,111 @@ module Aws
       end
 
       def abort_upload(upload_id, options, errors)
-        @client.abort_multipart_upload(
-          bucket: options[:bucket],
-          key: options[:key],
-          upload_id: upload_id
-        )
-        msg = "multipart upload failed: #{errors.map(&:message).join("; ")}"
+        @client.abort_multipart_upload(bucket: options[:bucket], key: options[:key], upload_id: upload_id)
+        msg = "multipart upload failed: #{errors.map(&:message).join('; ')}"
         raise MultipartUploadError.new(msg, errors)
-      rescue MultipartUploadError => error
-        raise error
-      rescue => error
-        msg = "failed to abort multipart upload: #{error.message}"
-        raise MultipartUploadError.new(msg, errors + [error])
+      rescue MultipartUploadError => e
+        raise e
+      rescue StandardError => e
+        msg = "failed to abort multipart upload: #{e.message}. " \
+              "Multipart upload failed: #{errors.map(&:message).join('; ')}"
+        raise MultipartUploadError.new(msg, errors + [e])
       end
 
-      def compute_parts(upload_id, source, options)
-        size = File.size(source)
-        default_part_size = compute_default_part_size(size)
+      def compute_parts(upload_id, source, file_size, options)
+        default_part_size = compute_default_part_size(file_size)
         offset = 0
         part_number = 1
         parts = []
-        while offset < size
-          parts << upload_part_opts(options).merge({
+        while offset < file_size
+          parts << upload_part_opts(options).merge(
             upload_id: upload_id,
             part_number: part_number,
             body: FilePart.new(
               source: source,
               offset: offset,
-              size: part_size(size, default_part_size, offset)
+              size: part_size(file_size, default_part_size, offset)
             )
-          })
+          )
           part_number += 1
           offset += default_part_size
         end
         parts
       end
 
+      def checksum_key?(key)
+        CHECKSUM_KEYS.include?(key)
+      end
+
+      def has_checksum_key?(keys)
+        keys.any? { |key| checksum_key?(key) }
+      end
+
+      def checksum_not_required?(options)
+        @client.config.request_checksum_calculation == 'when_required' && !options[:checksum_algorithm]
+      end
+
       def create_opts(options)
-        CREATE_OPTIONS.inject({}) do |hash, key|
-          hash[key] = options[key] if options.key?(key)
-          hash
+        opts = {}
+        unless checksum_not_required?(options)
+          opts[:checksum_algorithm] = Aws::Plugins::ChecksumAlgorithm::DEFAULT_CHECKSUM
         end
+        opts[:checksum_type] = 'FULL_OBJECT' if has_checksum_key?(options.keys)
+        CREATE_OPTIONS.each_with_object(opts) { |k, h| h[k] = options[k] if options.key?(k) }
+      end
+
+      def complete_opts(options)
+        opts = {}
+        opts[:checksum_type] = 'FULL_OBJECT' if has_checksum_key?(options.keys)
+        COMPLETE_OPTIONS.each_with_object(opts) { |k, h| h[k] = options[k] if options.key?(k) }
       end
 
       def upload_part_opts(options)
-        UPLOAD_PART_OPTIONS.inject({}) do |hash, key|
-          hash[key] = options[key] if options.key?(key)
-          hash
+        UPLOAD_PART_OPTIONS.each_with_object({}) do |key, hash|
+          # don't pass through checksum calculations
+          hash[key] = options[key] if options.key?(key) && !checksum_key?(key)
         end
       end
 
-      def upload_in_threads(pending, completed)
-        threads = []
-        @thread_count.times do
-          thread = Thread.new do
-            begin
-              while part = pending.shift
-                resp = @client.upload_part(part)
-                part[:body].close
-                completed.push(etag: resp.etag, part_number: part[:part_number])
-              end
-              nil
-            rescue => error
-              # keep other threads from uploading other parts
-              pending.clear!
-              error
-            end
+      def upload_with_executor(pending, completed, options)
+        upload_attempts = 0
+        completion_queue = Queue.new
+        abort_upload = false
+        errors = []
+        progress = MultipartProgress.new(pending, options[:progress_callback])
+
+        while (part = pending.shift)
+          break if abort_upload
+
+          upload_attempts += 1
+          @executor.post(part) do |p|
+            update_progress(progress, p)
+            resp = @client.upload_part(p)
+            p[:body].close
+            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
+            completion_queue << :done
           end
-          thread.abort_on_exception = true
-          threads << thread
         end
-        threads.map(&:value).compact
+
+        upload_attempts.times { completion_queue.pop }
+        errors
+      end
+
+      def apply_part_checksum(resp, part)
+        return unless (checksum = resp.context.params[:checksum_algorithm])
+
+        k = :"checksum_#{checksum.downcase}"
+        part[k] = resp.send(k)
       end
 
-      def compute_default_part_size(source_size)
-        [(source_size.to_f / MAX_PARTS).ceil, MIN_PART_SIZE].max.to_i
+      def compute_default_part_size(file_size)
+        [(file_size.to_f / MAX_PARTS).ceil, MIN_PART_SIZE].max.to_i
       end
 
       def part_size(total_size, part_size, offset)
@@ -157,9 +187,17 @@ module Aws
         end
       end
 
+      def update_progress(progress, part)
+        return unless progress.progress_callback
+
+        part[:on_chunk_sent] =
+          proc do |_chunk, bytes, _total|
+            progress.call(part[:part_number], bytes)
+          end
+      end
+
       # @api private
       class PartList
-
         def initialize(parts = [])
           @parts = parts
           @mutex = Mutex.new
@@ -177,10 +215,34 @@ module Aws
           @mutex.synchronize { @parts.clear }
         end
 
+        def size
+          @mutex.synchronize { @parts.size }
+        end
+
+        def part_sizes
+          @mutex.synchronize { @parts.map { |p| p[:body].size } }
+        end
+
         def to_a
           @mutex.synchronize { @parts.dup }
         end
+      end
+
+      # @api private
+      class MultipartProgress
+        def initialize(parts, progress_callback)
+          @bytes_sent = Array.new(parts.size, 0)
+          @total_sizes = parts.part_sizes
+          @progress_callback = progress_callback
+        end
 
+        attr_reader :progress_callback
+
+        def call(part_number, bytes_read)
+          # part numbers start at 1
+          @bytes_sent[part_number - 1] = bytes_read
+          @progress_callback.call(@bytes_sent, @total_sizes)
+        end
       end
     end
   end

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

@@ -0,0 +1,191 @@
+# frozen_string_literal: true
+
+require 'thread'
+require 'set'
+require 'tempfile'
+require 'stringio'
+
+module Aws
+  module S3
+    # @api private
+    class MultipartStreamUploader
+
+      DEFAULT_PART_SIZE = 5 * 1024 * 1024 # 5MB
+      CREATE_OPTIONS = Set.new(Client.api.operation(:create_multipart_upload).input.shape.member_names)
+      UPLOAD_PART_OPTIONS = Set.new(Client.api.operation(:upload_part).input.shape.member_names)
+      COMPLETE_UPLOAD_OPTIONS = Set.new(Client.api.operation(:complete_multipart_upload).input.shape.member_names)
+
+      # @option options [Client] :client
+      def initialize(options = {})
+        @client = options[:client] || Client.new
+        @executor = options[:executor]
+        @tempfile = options[:tempfile]
+        @part_size = options[:part_size] || DEFAULT_PART_SIZE
+      end
+
+      # @return [Client]
+      attr_reader :client
+
+      # @option options [required,String] :bucket
+      # @option options [required,String] :key
+      # @return [Seahorse::Client::Response] - the CompleteMultipartUploadResponse
+      def upload(options = {}, &block)
+        Aws::Plugins::UserAgent.metric('S3_TRANSFER') do
+          upload_id = initiate_upload(options)
+          parts = upload_parts(upload_id, options, &block)
+          complete_upload(upload_id, parts, options)
+        end
+      end
+
+      private
+
+      def initiate_upload(options)
+        @client.create_multipart_upload(create_opts(options)).upload_id
+      end
+
+      def complete_upload(upload_id, parts, options)
+        @client.complete_multipart_upload(
+          **complete_opts(options).merge(upload_id: upload_id, multipart_upload: { parts: parts })
+        )
+      rescue StandardError => e
+        abort_upload(upload_id, options, [e])
+      end
+
+      def upload_parts(upload_id, options, &block)
+        completed_parts = Queue.new
+        errors = []
+
+        begin
+          IO.pipe do |read_pipe, write_pipe|
+            upload_thread = Thread.new do
+              upload_with_executor(
+                read_pipe,
+                completed_parts,
+                errors,
+                upload_part_opts(options).merge(upload_id: upload_id)
+              )
+            end
+
+            block.call(write_pipe)
+          ensure
+            # Ensure the pipe is closed to avoid https://github.com/jruby/jruby/issues/6111
+            write_pipe.close
+            upload_thread.join
+          end
+        rescue StandardError => e
+          errors << e
+        end
+        return ordered_parts(completed_parts) if errors.empty?
+
+        abort_upload(upload_id, options, errors)
+      end
+
+      def abort_upload(upload_id, options, errors)
+        @client.abort_multipart_upload(bucket: options[:bucket], key: options[:key], upload_id: upload_id)
+        msg = "multipart upload failed: #{errors.map(&:message).join('; ')}"
+        raise MultipartUploadError.new(msg, errors)
+      rescue MultipartUploadError => e
+        raise e
+      rescue StandardError => e
+        msg = "failed to abort multipart upload: #{e.message}. "\
+          "Multipart upload failed: #{errors.map(&:message).join('; ')}"
+        raise MultipartUploadError.new(msg, errors + [e])
+      end
+
+      def create_opts(options)
+        CREATE_OPTIONS.each_with_object({}) do |key, hash|
+          hash[key] = options[key] if options.key?(key)
+        end
+      end
+
+      def upload_part_opts(options)
+        UPLOAD_PART_OPTIONS.each_with_object({}) do |key, hash|
+          hash[key] = options[key] if options.key?(key)
+        end
+      end
+
+      def complete_opts(options)
+        COMPLETE_UPLOAD_OPTIONS.each_with_object({}) do |key, hash|
+          hash[key] = options[key] if options.key?(key)
+        end
+      end
+
+      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)
+            temp_io.close
+            temp_io.unlink
+          end
+          nil
+        else
+          temp_io
+        end
+      end
+
+      def upload_with_executor(read_pipe, completed, errors, options)
+        completion_queue = Queue.new
+        queued_parts = 0
+        part_number = 0
+        mutex = Mutex.new
+        loop do
+          part_body, current_part_num = mutex.synchronize do
+            [read_to_part_body(read_pipe), part_number += 1]
+          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)
+          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
+        end
+        queued_parts.times { completion_queue.pop }
+      end
+
+      def create_completed_part(resp, part)
+        completed_part = { etag: resp.etag, part_number: part[:part_number] }
+        return completed_part unless part[:checksum_algorithm]
+
+        # get the requested checksum from the response
+        k = "checksum_#{part[:checksum_algorithm].downcase}".to_sym
+        completed_part[k] = resp[k]
+        completed_part
+      end
+
+      def ordered_parts(parts)
+        sorted = []
+        until parts.empty?
+          part = parts.pop
+          index = sorted.bsearch_index { |p| p[:part_number] >= part[:part_number] } || sorted.size
+          sorted.insert(index, part)
+        end
+        sorted
+      end
+
+      def clear_body(body)
+        if body.is_a?(Tempfile)
+          body.close
+          body.unlink
+        elsif body.is_a?(StringIO)
+          body.string.clear
+        end
+      end
+    end
+  end
+end